app.py

from functools import partial
import argparse
import cv2
import glob
import json
import logging
import math
import os
import sys
import re
import numpy as np
from typing import List, Optional
from PIL import Image, ImageFile
import tempfile
import datetime
import gradio as gr
import torch
from torch import nn
import torch.backends.cudnn as cudnn
from torchvision import transforms as pth_transforms
import shutil
import os
import spaces

os.system("wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth")
os.system("wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth")


sys.path.append("./segment-anything")
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor

sys.path.append(".")
from utils.data_utils import gen_square_crops
from facenet_pytorch import MTCNN, InceptionResnetV1

logger = logging.getLogger("dinov2")

# Default save paths for segmentation
OBJECT_SAVE_PATH = "./database/Objects/masks"
FACE_SAVE_PATH = "./database/Faces/masks"

# Initialize SAM model
def initialize_sam(sam_checkpoint, model_type="vit_h"):
    sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
    #sam.to(device="cuda" if torch.cuda.is_available() else "cpu")
    return sam

# Path to the SAM checkpoint
sam_checkpoint = "./sam_vit_h_4b8939.pth"
sam = initialize_sam(sam_checkpoint)
predictor = None


# Load RADIO model
model_version = "radio_v2.5-h"  # Using RADIOv2.5-H model (ViT-H/16)
model = torch.hub.load('NVlabs/RADIO', 'radio_model', version=model_version, progress=True, skip_validation=True)
#model.cuda().eval()

@spaces.GPU
def extract_features(image_path):
    model.cuda().eval()
    """Extract features from an image using the RADIO model."""
    x = Image.open(image_path).convert('RGB')
    x = pil_to_tensor(x).to(dtype=torch.float32, device='cuda')
    x.div_(255.0)  # RADIO expects values between 0 and 1
    x = x.unsqueeze(0)  # Add batch dimension
    
    # Resize to nearest supported resolution
    nearest_res = model.get_nearest_supported_resolution(*x.shape[-2:])
    x = F.interpolate(x, nearest_res, mode='bilinear', align_corners=False)
    
    # If using E-RADIO model, set optimal window size
    if "e-radio" in model_version:
        model.model.set_optimal_window_size(x.shape[2:])
    
    # Extract features - we're using the summary features for similarity comparison
    with torch.no_grad(), torch.autocast('cuda', dtype=torch.bfloat16):
        summary, spatial_feature = model(x)
    #spatial_featurev = spatial_feature.mean(dim=1)
    return summary

from torch.utils.data import DataLoader
from torchvision.transforms.functional import pil_to_tensor

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
from PIL import Image
from tqdm import tqdm # For the progress bar
# import numpy as np # Required if your pil_to_tensor function uses it (e.g., for np.array)

# ------------- BEGIN ASSUMPTIONS / REQUIRED EXTERNAL DEFINITIONS ---------------
# The following variables/functions (`model`, `model_version`, `pil_to_tensor`)
# are assumed to be defined and accessible in the scope where `extract_features`
# is called. For example, they could be global variables, or `extract_features`
# could be a method of a class that holds them as attributes (e.g., `self.model`).

# 1. `model`: A PyTorch model object (expected to be on CUDA).
#    - Must have a method `get_nearest_supported_resolution(height, width)` which
#      returns a tuple (new_height, new_width).
#    - If `model_version` (see below) contains "e-radio", `model.model` (or the relevant
#      submodule) must have a method `set_optimal_window_size((height, width))`.
#    - The model's forward pass should accept a batch of tensors `(B, C, H, W)` and
#      return a tuple `(summary_batch, spatial_feature_batch)`.

# 2. `model_version`: A string indicating the model version (e.g., "e-radio_v1.0").
#    This is used to conditionally call `set_optimal_window_size`.

# 3. `pil_to_tensor`: A function that converts a PIL Image object to a PyTorch tensor.
#    - Input: A PIL Image object (typically after `.convert('RGB')`).
#    - Output: A PyTorch tensor, expected to be in CHW (Channels, Height, Width) format.
#    - IMPORTANT: Based on the original code snippet:
#      `x = pil_to_tensor(x).to(dtype=torch.float32, device='cuda')`
#      `x.div_(255.0)`
#      This sequence implies that `pil_to_tensor(x)` returns a tensor with pixel values
#      in the range [0, 255] (e.g., a `torch.ByteTensor` or a `torch.FloatTensor`
#      representing unnormalized pixel values). It should NOT normalize the tensor to
#      the [0, 1] range itself, as `div_(255.0)` handles this.

# Example placeholder (ensure these are correctly defined in your actual environment):
# def example_pil_to_tensor(pil_image):
#     import numpy as np
#     return torch.as_tensor(np.array(pil_image)).permute(2, 0, 1)
# pil_to_tensor = example_pil_to_tensor
# class DummyModel(torch.nn.Module): # Replace with your actual model
#     def __init__(self): super().__init__(); self.model = self
#     def get_nearest_supported_resolution(self, h, w): return h, w
#     def set_optimal_window_size(self, hw): pass
#     def forward(self, x): return torch.rand(x.shape[0], 10, device=x.device), None
# model = DummyModel().to('cuda')
# model_version = "e-radio_test"
# ------------- END ASSUMPTIONS / REQUIRED EXTERNAL DEFINITIONS ---------------

def _robust_collate_fn_for_extract_features(batch):
    """
    Custom collate_fn for DataLoader. Batches indices using default_collate
    and returns image data (paths, PIL.Images, or torch.Tensors) as a list.
    """
    image_data_list = [item[0] for item in batch]
    indices = [item[1] for item in batch]
    batched_indices = torch.utils.data.default_collate(indices)
    return image_data_list, batched_indices

@spaces.GPU
def extract_features(object_dataset, batch_size, num_workers):
     """
     Extracts features from images, handling inputs as paths, PIL Images, or Tensors.
     Assumes `model`, `model_version`, `pil_to_tensor` are in calling scope.
     """
     model.cuda().eval()
     dataloader = DataLoader(
         object_dataset,
         batch_size=batch_size,
         num_workers=num_workers,
         shuffle=False,
         collate_fn=_robust_collate_fn_for_extract_features
     )

     all_summaries = []
     pre_resize_target_h_w = (256, 256)
     if hasattr(object_dataset, 'imsize') and object_dataset.imsize is not None:
         if isinstance(object_dataset.imsize, int):
             pre_resize_target_h_w = (object_dataset.imsize, object_dataset.imsize)
         elif isinstance(object_dataset.imsize, (list, tuple)) and len(object_dataset.imsize) == 2:
             pre_resize_target_h_w = tuple(map(int, object_dataset.imsize))
         else:
             print(f"Warning: `object_dataset.imsize` format ({object_dataset.imsize}) "
                   f"is not recognized. Using default pre-resize: {pre_resize_target_h_w}.")

     for batch_of_image_data, _ in tqdm(dataloader, desc="Extracting Features"):
         current_batch_processed_tensors = []
         for image_data_item in batch_of_image_data:
             x = None # Initialize x, which will become the processed tensor

             if isinstance(image_data_item, str):  # Item is an image path
                 img_pil = Image.open(image_data_item).convert('RGB')
                 x = pil_to_tensor(img_pil) # Expected CHW, [0,255] range (any dtype)
                 x = x.to(dtype=torch.float32, device='cuda')
                 x.div_(255.0) # Normalize to [0,1]
             elif isinstance(image_data_item, Image.Image):  # Item is already a PIL Image
                 img_pil = image_data_item.convert('RGB') # Ensure RGB
                 x = pil_to_tensor(img_pil) # Expected CHW, [0,255] range (any dtype)
                 x = x.to(dtype=torch.float32, device='cuda')
                 x.div_(255.0) # Normalize to [0,1]
             elif isinstance(image_data_item, torch.Tensor):  # Item is a PyTorch Tensor
                 # Assume the input tensor also needs to be processed like the output of pil_to_tensor:
                 # i.e., converted to float32, moved to cuda, and then normalized from a [0,255] scale to [0,1].
                 # This is a strong assumption; if dataset tensors are already [0,1] float, this div_ is wrong.
                 # However, it aligns with the normalization applied in other branches.
                 x = image_data_item.to(dtype=torch.float32, device='cuda')
                 # If the input tensor might already be in [0,1] range:
                 if x.max() > 1.1: # Heuristic: if max value suggests it's not [0,1]
                     # This warning is helpful to understand if assumptions are met.
                     # print(f"Note: Input tensor (max val: {x.max().item():.2f}) "
                     #       f"appears to be in [0,255] range. Normalizing by dividing by 255.")
                     x.div_(255.0)
                 elif not (0 <= x.min() and x.max() <= 1.1): # check if it's not in a typical normalized range or close to it.
                      # Handle cases like negative values or values slightly outside expected bounds for [0,1] images.
                      # If it is not in [0,1] and not clearly in [0,255] (e.g. [-1,1]), this path might need more specific logic.
                      # For now, if it's not clearly [0,255] scaled (max > 1.1), we assume it is either already [0,1]
                      # or requires a different normalization not covered here.
                      # The current logic implicitly assumes floats not >1.1 are already okay.
                      pass # Assume float tensors with max <= 1.1 are already normalized or don't need div by 255.
             else:
                 raise TypeError(
                     f"Dataset provided an item of unexpected type for image data: {type(image_data_item)}. "
                     f"Expected a path string, a PIL.Image object, or a torch.Tensor."
                 )
            
             # Common processing for x (now a CUDA float32 tensor, intended to be [0,1])
             if x.shape[1:] != pre_resize_target_h_w:
                 x = F.interpolate(x.unsqueeze(0),
                                   size=pre_resize_target_h_w,
                                   mode='bilinear',
                                   align_corners=False).squeeze(0)

             current_batch_processed_tensors.append(x)

         if not current_batch_processed_tensors:
             continue

         x_batch = torch.stack(current_batch_processed_tensors)
         nearest_res = model.get_nearest_supported_resolution(*x_batch.shape[-2:])
         x_batch = F.interpolate(x_batch, nearest_res, mode='bilinear', align_corners=False)

         if "e-radio" in model_version:
             target_module_for_window_size = None
             if hasattr(model, 'model') and hasattr(model.model, 'set_optimal_window_size'):
                 target_module_for_window_size = model.model
             elif hasattr(model, 'set_optimal_window_size'):
                 target_module_for_window_size = model
            
             if target_module_for_window_size:
                 target_module_for_window_size.set_optimal_window_size(x_batch.shape[2:])
             else:
                 print(f"Warning: 'e-radio' in model_version, but 'set_optimal_window_size' method not found.")

         with torch.no_grad(), torch.autocast('cuda', dtype=torch.bfloat16):
             summary_batch, _ = model(x_batch)

         all_summaries.append(summary_batch)

     """if not all_summaries:
         return torch.empty(0, device='cpu'),
     """
     final_summaries = torch.cat(all_summaries, dim=0)
     return final_summaries

import re

def remove_numbers_in_brackets(text):
    """
    移除字符串中所有格式为[数字]的内容
    
    参数:
        text (str): 需要处理的字符串
    
    返回:
        str: 处理后的字符串
    """
    # 使用正则表达式匹配[数字]模式并替换为空字符串
    # \[ 匹配左方括号
    # \d+ 匹配一个或多个数字
    # \] 匹配右方括号
    return re.sub(r'\[\d+\]', '', text)



# Global state to track masks and image information
class AppState:
    def __init__(self):
        self.current_image_index = 0
        self.images = []  # List of (image_array, image_name)
        self.masks = []   # List of masks corresponding to images
        self.gallery_items = []  # List of (image, caption) for gallery
        self.current_object_id = None  # Track the current object ID
        self.processed_count = 0  # Counter for processed objects
        self.object_image_counts = {}  # Counter for images per object
        self.mode = "object"  # Default mode is object segmentation
        self.reset()
    
    def reset(self):
        self.current_image_index = 0
        self.images = []
        self.masks = [None] * 100  # Pre-allocate for potential uploads
        # Don't reset gallery, object ID, counters, or mode
    
    def add_images(self, image_list):
        """添加图像到状态中"""
        self.reset()
        for img in image_list:
            if img is not None:
                # 确保图像是正确的格式（RGB numpy 数组）
                if isinstance(img, str):
                    # 这是一个文件路径
                    try:
                        img_array = np.array(Image.open(img).convert('RGB'))
                        img_name = os.path.basename(img)
                        self.images.append((img_array, img_name))
                    except Exception as e:
                        print(f"Error loading image {img}: {str(e)}")
                elif hasattr(img, 'name'):
                    # 这是一个带有 name 属性的类文件对象
                    try:
                        img_array = np.array(Image.open(img.name).convert('RGB'))
                        img_name = os.path.basename(img.name)
                        self.images.append((img_array, img_name))
                    except Exception as e:
                        print(f"Error loading image object: {str(e)}")
                else:
                    # 这可能已经是一个图像数组
                    try:
                        if isinstance(img, Image.Image):
                            # PIL 图像对象
                            img_array = np.array(img.convert('RGB'))
                        else:
                            # 假设是 numpy 数组
                            img_array = np.array(img)
                        img_name = f"image_{len(self.images)}.png"
                        self.images.append((img_array, img_name))
                    except Exception as e:
                        print(f"Error processing image data: {str(e)}")
        return len(self.images)
    
    def get_current_image(self):
        if 0 <= self.current_image_index < len(self.images):
            return self.images[self.current_image_index][0]
        return None
    
    def get_current_image_name(self):
        if 0 <= self.current_image_index < len(self.images):
            return self.images[self.current_image_index][1]
        return None
    
    def set_mask(self, mask):
        if 0 <= self.current_image_index < len(self.images):
            self.masks[self.current_image_index] = mask
    
    def get_current_mask(self):
        if 0 <= self.current_image_index < len(self.images):
            return self.masks[self.current_image_index]
        return None
    
    def next_image(self):
        if len(self.images) > 0:
            self.current_image_index = (self.current_image_index + 1) % len(self.images)
        return self.current_image_index
    
    def get_status_text(self):
        if len(self.images) == 0:
            return "No images loaded"
        
        item_type = "Face" if self.mode == "👤face" else "Object"
        item_text = f"{item_type} ID: {self.current_object_id}" if self.current_object_id else f"New {item_type}"
        return f"Image {self.current_image_index + 1}/{len(self.images)}: {self.get_current_image_name()} | {item_text}"
    
    def add_to_gallery(self, image, caption):
        """Add an image and its caption to the gallery"""
        self.gallery_items.append((image, caption))
        return self.gallery_items
    
    def get_gallery(self):
        """Return the gallery items in the format needed for gr.Gallery"""
        return self.gallery_items
    
    def get_next_object_id(self):
        """Get the next object ID for a new object"""
        self.processed_count += 1
        self.current_object_id = f"{self.processed_count:03d}"
        self.object_image_counts[self.current_object_id] = 0
        return self.current_object_id
    
    def get_next_image_id(self):
        """Get the next image ID for the current object"""
        if self.current_object_id is None:
            self.get_next_object_id()
        
        self.object_image_counts[self.current_object_id] += 1
        return f"{self.object_image_counts[self.current_object_id]:03d}"

# Create state for segmentation module
state = AppState()

# Function to update mode
def update_mode(new_mode):
    state.mode = new_mode
    # Reset object ID when changing modes
    state.current_object_id = None
    return f"Mode changed to: {new_mode.capitalize()} segmentation"

def create_masked_object(image, mask, margin=0):
    """
    Create a masked image with white background, cropped to the object plus a margin.
    
    Args:
        image: Original image (numpy array)
        mask: Binary mask (numpy array, same size as image)
        margin: Number of pixels to add around the object bounding box
    
    Returns:
        Masked image with white background, cropped to the object
    """
    # Find the bounding box of the object in the mask
    y_indices, x_indices = np.where(mask)
    if len(y_indices) == 0 or len(x_indices) == 0:
        return image  # If mask is empty, return original image
    
    # Get bounding box coordinates with margin
    y_min, y_max = max(0, np.min(y_indices) - margin), min(image.shape[0], np.max(y_indices) + margin)
    x_min, x_max = max(0, np.min(x_indices) - margin), min(image.shape[1], np.max(x_indices) + margin)
    
    # Create a white background image of the cropped size
    cropped_size = (y_max - y_min, x_max - x_min, 3)
    masked_image = np.ones(cropped_size, dtype=np.uint8) * 255
    
    # Copy the object pixels from the original image
    mask_cropped = mask[y_min:y_max, x_min:x_max]
    masked_image[mask_cropped] = image[y_min:y_max, x_min:x_max][mask_cropped]
    
    return masked_image

def upload_images(image_list):
    """处理图像上传，存储所有图像，并返回第一个图像"""
    count = state.add_images(image_list)
    
    if count == 0:
        return None, None, f"No valid images uploaded", state.get_gallery()
    
    current_image = state.get_current_image()
    return current_image, None, f"Uploaded {count} images. Viewing image 1/{count}: {state.get_current_image_name()}", state.get_gallery()

def handle_example_selection(mode, file_paths, file_output, object_info):
    """处理从示例中选择图像的事件"""
    # 更新模式
    state.mode = mode
    
    # 确保路径是列表
    if isinstance(file_paths, str):
        file_paths = [file_paths]
    
    # 处理图像上传
    count = state.add_images(file_paths)
    
    if count == 0:
        return None, None, f"No valid images uploaded", state.get_gallery(), object_info
    
    current_image = state.get_current_image()
    status = f"Loaded example image 1/{count}: {state.get_current_image_name()}"
    
    return current_image, None, status, state.get_gallery(), object_info

def navigate_images(is_same_object=False):
    """Navigate to the next image"""
    # If it's not the same object, reset the object ID
    if not is_same_object:
        state.current_object_id = None  # This will trigger a new object ID on save
    
    state.next_image()
    
    current_image = state.get_current_image()
    if current_image is None:
        return None, None, "No images available", state.get_gallery(), None  # Return None to clear file upload
    
    # Get mask if previously generated
    current_mask = state.get_current_mask()
    mask_display = None
    
    if current_mask is not None:
        # Create visual representation of mask
        img_cv = state.get_current_image()
        colored_mask = np.zeros_like(img_cv)
        colored_mask[current_mask] = [0, 0, 255]  # Red mask
        blended = cv2.addWeighted(img_cv, 0.7, colored_mask, 0.3, 0)
        mask_display = blended
    
    status_text = state.get_status_text()
    return current_image, mask_display, status_text, state.get_gallery(), None  # Return None to clear file upload

    
@spaces.GPU
def generate_mask(image, evt: gr.SelectData): # 'image' is the numpy array from the clicked component
    sam.to(device="cuda" if torch.cuda.is_available() else "cpu")
    global predictor

    # Use the image passed by the event!
    if image is None:
        return None, None, "Cannot segment: Image component is empty.", state.get_gallery()

    # Ensure the image is a NumPy array in RGB format (Gradio usually provides this)
    if not isinstance(image, np.ndarray):
        try:
            # Attempt conversion if needed (e.g., if PIL Image was somehow passed)
            image_cv = np.array(Image.fromarray(image).convert('RGB'))
        except Exception as e:
            print(f"Warning: Could not convert input image for segmentation: {e}")
            # Fallback to state as a last resort, or return error
            image_from_state = state.get_current_image()
            if image_from_state is None:
                 return None, None, f"Error processing image data and state unavailable.", state.get_gallery()
            image_cv = image_from_state # Use state image if event image fails
    else:
        image_cv = image # Already a numpy array

    # Ensure 3 channels (RGB)
    if len(image_cv.shape) == 2: # Grayscale
        image_cv = cv2.cvtColor(image_cv, cv2.COLOR_GRAY2RGB)
    elif image_cv.shape[2] == 4: # RGBA
        image_cv = cv2.cvtColor(image_cv, cv2.COLOR_RGBA2RGB)
    elif image_cv.shape[2] != 3:
         return None, None, f"Unsupported image format: channels={image_cv.shape[2]}", state.get_gallery()

    # Check state only for context (like image name) and storing the mask later
    current_image_name = "unknown_image"
    current_image_state_valid = (0 <= state.current_image_index < len(state.images))
    if current_image_state_valid:
        current_image_name = state.get_current_image_name()
    else:
        print("Warning: State index out of bounds, using default name.")


    # Initialize the predictor
    try:
        predictor = SamPredictor(sam)
        predictor.set_image(image_cv) # Use the image_cv derived *directly* from the event argument!
    except Exception as e:
        print(f"Error setting image in SAM predictor: {e}")
        return None, None, f"Error setting image in SAM predictor: {e}", state.get_gallery()

    # Get coordinates from the click event
    input_point = np.array([[evt.index[0], evt.index[1]]])
    input_label = np.array([1])

    # Generate masks
    try:
        masks, scores, logits = predictor.predict(
            point_coords=input_point,
            point_labels=input_label,
            multimask_output=True,
        )
    except Exception as e:
        print(f"Error during SAM prediction: {e}")
        return None, None, f"Error during SAM prediction: {e}", state.get_gallery()

    if masks is None or len(masks) == 0:
         return None, None, "SAM prediction failed to produce masks.", state.get_gallery()

    # Get the best mask
    mask_idx = np.argmax(scores)
    mask = masks[mask_idx]
    score = scores[mask_idx]

    # --- Storing the mask ---
    # We still need the state to know *where* to store this mask.
    if current_image_state_valid:
        state.set_mask(mask) # Store mask in state corresponding to the current index
    else:
        print("Warning: Could not store mask in state due to invalid index.")

    # Apply mask visualization to the image_cv we processed
    colored_mask = np.zeros_like(image_cv)
    colored_mask[mask] = [0, 0, 255] # Red mask (or choose another color)
    blended = cv2.addWeighted(image_cv, 0.7, colored_mask, 0.3, 0)

    # Create binary mask image for potential display or use (optional)
    binary_mask_img = Image.fromarray((mask * 255).astype(np.uint8))

    status_msg = f"Generated mask for {current_image_name} (Score: {score:.2f})"

    # Return the blended image to show the mask, and the status.
    # We return blended (visual) and binary_mask_img (data) potentially for different outputs if needed.
    # For your current setup, you might just need the blended image in the 'masked_image' output.
    # The second output component 'masked_image' currently expects the binary mask, let's adjust return slightly
    # If masked_image component should show the blended image:
    # return blended, status_msg, state.get_gallery()
    # If masked_image component should show the pure binary mask:
    # return blended, binary_mask_img, status_msg, state.get_gallery()
    # Your original code returned blended -> masked_image[0], binary_mask -> masked_image[1], status -> status_text, gallery -> gallery
    # Let's match that structure assuming masked_image was intended to potentially show both forms or the binary form
    return binary_mask_img, status_msg, state.get_gallery() # Match original output count/types roughly

def save_mask_and_text(object_name):
    current_mask = state.get_current_mask()
    current_image = state.get_current_image()
    
    # Set save path based on current mode
    save_path = FACE_SAVE_PATH if state.mode == "👤face" else OBJECT_SAVE_PATH
    
    if current_mask is None:
        return f"No mask has been generated yet. Please click on an {state.mode} first.", state.get_gallery()
    
    # Get or create object ID
    if state.current_object_id is None:
        # Check existing directories to determine the next ID
        existing_dirs = sorted(glob.glob(os.path.join(save_path, '*')))
        if existing_dirs:
            # Extract existing IDs and find the highest
            existing_ids = []
            for dir_path in existing_dirs:
                try:
                    # Try to extract numeric ID from directory name
                    dir_id = os.path.basename(dir_path)
                    if dir_id.isdigit():
                        existing_ids.append(int(dir_id))
                except ValueError:
                    continue
            
            # Get the next ID in sequence if any exist
            if existing_ids:
                next_id = max(existing_ids) + 1
                object_id = f"{next_id:03d}"
            else:
                # Start from 001 if no valid numeric IDs found
                object_id = "001"
        else:
            # No existing directories, start from 001
            object_id = "001"
        
        # Set the state's current object ID
        state.current_object_id = object_id
    else:
        object_id = state.current_object_id
    
    # Create object-specific directory
    object_dir = os.path.join(save_path, object_id)
    os.makedirs(object_dir, exist_ok=True)
    
    # Get the next image ID for this specific object
    # Check existing image files to determine the next image ID
    image_dir = os.path.join(object_dir, 'images')
    os.makedirs(image_dir, exist_ok=True)
    
    existing_images = sorted(
        glob.glob(os.path.join(image_dir, '*.png')) +
        glob.glob(os.path.join(image_dir, '*.jpg')) +
        glob.glob(os.path.join(image_dir, '*.jpeg')) +
        glob.glob(os.path.join(image_dir, '*.bmp'))
    )
    
    if existing_images:
        # Extract existing image IDs and find the highest
        existing_img_ids = []
        for img_path in existing_images:
            try:
                # Extract numeric ID from filename (without extension)
                img_id = os.path.splitext(os.path.basename(img_path))[0]
                if img_id.isdigit():
                    existing_img_ids.append(int(img_id))
            except ValueError:
                continue
        
        # Get the next image ID in sequence if any exist
        if existing_img_ids:
            next_img_id = max(existing_img_ids) + 1
            image_id = f"{next_img_id:03d}"
        else:
            # Start from 001 if no valid numeric IDs found
            image_id = "001"
    else:
        # No existing images, start from 001
        image_id = "001"
    
    # Generate filenames with object and image IDs
    current_image_name = state.get_current_image_name()
    image_stem = os.path.splitext(current_image_name)[0]
    
    # Define paths for all files we'll save
    mask_dir = os.path.join(object_dir, 'masks')
    os.makedirs(mask_dir, exist_ok=True)
    mask_path = os.path.join(mask_dir, f"{image_id}.png")
    
    image_path = os.path.join(image_dir, f"{image_id}.png")

    ann_dir = os.path.join(object_dir, 'anns')
    os.makedirs(ann_dir, exist_ok=True)
    text_path = os.path.join(ann_dir, f"{image_id}.txt")
    
    # Save binary mask
    mask_img = Image.fromarray((current_mask * 255).astype(np.uint8))
    mask_img.save(mask_path)
    
    # Create and save the masked object with white background
    masked_object = create_masked_object(current_image, current_mask)
    masked_obj_img = Image.fromarray(masked_object)
    masked_obj_img.save(image_path)
    object_name = object_name.replace('\n', '')
    
    # Save text information
    with open(text_path, 'w') as f:
        f.write(f"Object ID: {object_id}\n")
        f.write(f"Image Number: {image_id}\n")
        f.write(f"Object Name: {object_name}\n")
        f.write(f"Source Image: {current_image_name}\n")
        f.write(f"Creation Time: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
        f.write(f"Type: {state.mode}\n")  # Add type information
    
    # Add the masked object to the gallery with caption
    mode_prefix = "F" if state.mode == "👤face" else "O"
    caption = f"{mode_prefix}#{object_id}/{image_id}: {object_name} ({image_stem})"
    state.add_to_gallery(masked_object, caption)
    
    # Update status message
    status_msg = f"Saved {state.mode} {object_id}/{image_id}: {object_name}"
    
    return status_msg, state.get_gallery()

# ====== DETECTION PART ======

class RealWorldDataset(torch.utils.data.Dataset):
    def __init__(self, data_dir, dataset, data=None, transform=None, imsize=None):
        if dataset == 'Object':
            num_obj = []
            image_dir = []
            mask_dir = []
            count = []
            anns_dir = []
            source_list = sorted(glob.glob(os.path.join(data_dir, '*')))

            for _, source_dir in enumerate(source_list):
                num_obj.append(source_dir.split('/')[-1].split('.')[0])
                image_paths = sorted([p for p in glob.glob(os.path.join(source_dir, 'images', '*'))
                                      if re.search('/*\.(jpg|jpeg|png|gif|bmp|pbm)', str(p))])
                image_dir.extend(image_paths)
                mask_paths = sorted([p for p in glob.glob(os.path.join(source_dir, 'masks', '*'))
                                     if re.search('/*\.(jpg|jpeg|png|gif|bmp|pbm)', str(p))])
                mask_dir.extend(mask_paths)
                ann_paths = sorted([p for p in glob.glob(os.path.join(source_dir, 'anns', '*'))
                                     if re.search('/*\.(txt)', str(p))])
                anns_dir.append(ann_paths)
                count.append(len(image_paths))

            cfg = dict()
            cfg['dataset'] = dataset
            cfg['data_dir'] = data_dir
            cfg['image_dir'] = image_dir
            cfg['mask_dir'] = mask_dir
            cfg['obj_name'] = num_obj  # object lists for Object
            cfg['length'] = count
            cfg['anns_dir'] = anns_dir

            self.samples = cfg['image_dir']

        elif dataset == 'Scene':
            num_scene = []
            image_dir = []
            proposals = []
            count = []

            with open(os.path.join(os.path.dirname(data_dir),
                                   'proposals_on_' + data_dir.split('/')[-1] + '.json')) as f:
                proposal_json = json.load(f)

            source_list = sorted(glob.glob(os.path.join(data_dir, '*')))

            for idx, source_dir in enumerate(source_list):
                scene_name = source_dir.split('/')[-1]
                num_scene.append(scene_name)

                image_paths = sorted([p for p in glob.glob(os.path.join(source_dir, '*'))
                                      if re.search('/*\.(jpg|jpeg|png|gif|bmp|pbm)', str(p))])
                image_dir.extend(image_paths)
                count.append(len(image_paths))
                proposals.extend(proposal_json[scene_name])

            cfg = dict()
            cfg['dataset'] = dataset
            cfg['data_dir'] = data_dir
            cfg['image_dir'] = image_dir
            cfg['proposals'] = proposals
            cfg['scene_name'] = num_scene  # scene list for Scene
            cfg['length'] = count

            self.samples = cfg['image_dir']

        else:  # for demo scene image
            with open(os.path.join(data_dir, 'proposals_on_' + dataset + '.json')) as f:
                proposal_json = json.load(f)

            cfg = dict()
            cfg['dataset'] = dataset
            cfg['data_dir'] = data_dir
            cfg['image_dir'] = None
            cfg['proposals'] = proposal_json
            cfg['scene_name'] = [dataset]  # scene list for Scene
            cfg['length'] = [len(data)]

            self.samples = data

        self.cfg = cfg
        self.transform = transform
        self.imsize = imsize

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, index):
        if "test" in self.cfg['dataset']:
            img = self.samples[index]
        else:
            path = self.samples[index]
            ImageFile.LOAD_TRUNCATED_IMAGES = True

            with open(path, 'rb') as f:
                img = Image.open(f)
                img = img.convert('RGB')
        
        w, h = img.size

        if (self.imsize is not None) and (min(w, h) > self.imsize):
            img.thumbnail((self.imsize, self.imsize), Image.Resampling.LANCZOS)
            w, h = img.size 
        new_w = math.ceil(w / 14) * 14
        new_h = math.ceil(h / 14) * 14
        img = img.resize((new_w, new_h), Image.Resampling.LANCZOS)

        if self.transform is not None:
            img = self.transform(img)
        return img, index

def compute_similarity(obj_feats, roi_feats):
    """
    Compute Cosine similarity between object features and proposal features
    """
    roi_feats = roi_feats.unsqueeze(-2)
    sim = torch.nn.functional.cosine_similarity(roi_feats, obj_feats, dim=-1)
    return sim

def stableMatching(preferenceMat):
    """
    Compute Stable Matching
    """
    mDict = dict()

    engageMatrix = np.zeros_like(preferenceMat)
    for i in range(preferenceMat.shape[0]):
        tmp = preferenceMat[i]
        sortIndices = np.argsort(tmp)[::-1]
        mDict[i] = sortIndices.tolist()

    freeManList = list(range(preferenceMat.shape[0]))

    while freeManList:
        curMan = freeManList.pop(0)

        curWoman = mDict[curMan].pop(0)
        if engageMatrix[:, curWoman].sum() == 0:
            engageMatrix[curMan, curWoman] = 1
        else:
            engagedMan = np.where(engageMatrix[:, curWoman] == 1)[0][0]
            if preferenceMat[engagedMan, curWoman] > preferenceMat[curMan, curWoman]:
                freeManList.append(curMan)
            else:
                engageMatrix[engagedMan, curWoman] = 0
                engageMatrix[curMan, curWoman] = 1
                freeManList.append(engagedMan)
    return engageMatrix

def get_args_parser(
        description: Optional[str] = None,
        parents: Optional[List[argparse.ArgumentParser]] = [],
        add_help: bool = True,
):
    #setup_args_parser = get_setup_args_parser(parents=parents, add_help=False)
    #parents = [setup_args_parser]

    parser = argparse.ArgumentParser(
        description=description,
        parents=parents,
        add_help=add_help,
    )
    parser.add_argument(
        "--train_path",
        default="/mnt/14T-disk/code/Contextual_Referring_Understanding/OSLD/logo-images-split-by-company",
        type=str,
        help="Path to train dataset.",
    )
    parser.add_argument(
        "--test_path",
        default="/mnt/14T-disk/code/instance-detection/database/test",
        type=str,
        help="Path to test dataset.",
    )
    parser.add_argument(
        "--imsize",
        default=224,
        type=int,
        help="Image size",
    )
    parser.add_argument(
        "--pretrained_weights",
        default="dinov2_vitl14_pretrain.pth",
        type=str,
        help="Path to pretrained weights to evaluate.",
    )
    parser.add_argument(
        "--output_dir",
        default="./output",
        type=str,
        help="Path to save outputs.")
    parser.add_argument("--num_workers", default=0, type=int, help="Number of data loading workers per GPU.")

    parser.add_argument(
        "--gather-on-cpu",
        action="store_true",
        help="Whether to gather the train features on cpu, slower"
             "but useful to avoid OOM for large datasets (e.g. ImageNet22k).",
    )

    parser.set_defaults(
        train_dataset="Object",
        test_dataset="Scene",
        batch_size=1,
        num_workers=0,
    )
    return parser

def visualize_detection(image, results, object_names):
    """Visualize detection results on image"""
    output_img = image.copy()
    for i, res in enumerate(results):
        x, y, w, h = res['bbox']
        category = object_names[res['category_id']]
        score = res['score']
        
        # Convert to absolute coordinates based on scale
        x = int(x * res['scale'])
        y = int(y * res['scale'])
        w = int(w * res['scale'])
        h = int(h * res['scale'])
        
        # Draw rectangle
        cv2.rectangle(output_img, (x, y), (x+w, y+h), (0, 255, 0), 2)
        
        # Add label
        text = f"[{i}]: {score:.2f}"
        cv2.putText(output_img, text, (x, y-5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
    
    return output_img

# Initialize args globally
args_parser = get_args_parser(description="SAM-DINOv2 Instance Detection")
imsize = 224

args = args_parser.parse_args(args=[
    "--train_path", "./database/Objects/masks",
    "--test_path", "temp_path_placeholder",  # This will be updated during runtime
    "--pretrained_weights", "./dinov2_vitl14_reg4_pretrain.pth",
    "--output_dir", f"exps/output_RankSelect_{imsize}_mask",  # Default tag, will be updated
])

# Set up output directory and model once
os.makedirs(args.output_dir, exist_ok=True)
#model, autocast_dtype = setup_and_build_model(args)

def detect_objects(input_img, score_threshold=0.52, tag="mask"):
    """Main function to detect objects in an image"""
    # Create temporary file for the input image
    with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as f:
        temp_path = f.name
        input_img.save(temp_path)
    
    # Load object features
    transform = pth_transforms.Compose([pth_transforms.ToTensor(),])
    object_dataset = RealWorldDataset(args.train_path, args.train_dataset, transform=transform, imsize=args.imsize)
    if len(object_dataset) == 0:
        raw_image = np.array(input_img.convert('RGB'))
        return [], raw_image, []
    # Initialize variables for features
    need_extract_all = True
    existing_features = None
    dataset_size = len(object_dataset)
    
    # Check if features file exists
    if os.path.exists(os.path.join('./database/Objects', 'object_features.json')):
        with open(os.path.join('./database/Objects', 'object_features.json'), 'r') as f:
            feat_dict = json.load(f)
        
        # Check if dimensions match
        if 'features' in feat_dict and len(feat_dict['features']) == dataset_size:
            # All features already extracted
            object_features = torch.Tensor(feat_dict['features']).cuda()
            need_extract_all = False
        elif 'features' in feat_dict and len(feat_dict['features']) > 0:
            # Partial features exist
            existing_features = torch.Tensor(feat_dict['features']).cuda()
            print(f"Found {len(feat_dict['features'])} existing features, but dataset has {dataset_size} objects.")
            print(f"Will extract features for the remaining {dataset_size - len(feat_dict['features'])} objects.")
            need_extract_all = False
    
    if need_extract_all:
        # Extract features for all objects
        print("Extracting features for all objects...")
        object_features = extract_features(
             object_dataset, args.batch_size, args.num_workers
        )
        feat_dict = dict()
        feat_dict['features'] = object_features.detach().cpu().tolist()
        with open(os.path.join('./database/Objects', 'object_features.json'), 'w') as f:
            json.dump(feat_dict, f)
    elif existing_features is not None:
        # Create a subset dataset for unprocessed objects
        num_existing = existing_features.size(0)
        remaining_indices = list(range(num_existing, dataset_size))
        
        class SubsetDataset(torch.utils.data.Dataset):
            def __init__(self, dataset, indices):
                self.dataset = dataset
                self.indices = indices
                
            def __getitem__(self, idx):
                return self.dataset[self.indices[idx]]
                
            def __len__(self):
                return len(self.indices)
        
        remaining_dataset = SubsetDataset(object_dataset, remaining_indices)
        
        # Extract features for remaining objects
        print(f"Extracting features for {len(remaining_dataset)} remaining objects...")
        new_features = extract_features(
            remaining_dataset, args.batch_size, args.num_workers
        )
        
        # Combine existing and new features
        
        object_features = torch.cat([existing_features, new_features], dim=0)
        
        # Save the combined features
        feat_dict = dict()
        feat_dict['features'] = object_features.detach().cpu().tolist()
        with open(os.path.join('./database/Objects', 'object_features.json'), 'w') as f:
            json.dump(feat_dict, f)
    
    # Normalize features
    object_features = nn.functional.normalize(object_features, dim=1, p=2)
    
    # Generate masks using SAM
    raw_image = np.array(input_img.convert('RGB'))
    ratio = 0.25
    scene_image = cv2.resize(raw_image, (int(raw_image.shape[1] * ratio), int(raw_image.shape[0] * ratio)), cv2.INTER_LINEAR)
    
    mask_generator = SamAutomaticMaskGenerator(sam)
    masks = mask_generator.generate(scene_image)
    
    # Process masks to generate proposals
    image_height, image_width = raw_image.shape[:-1]
    scene_name = "test_scene"
    
    rois = []
    sel_rois = []
    for ind, segment_dict in enumerate(masks):
        # Get bbox
        x0 = int(segment_dict['bbox'][0])
        y0 = int(segment_dict['bbox'][1])
        x1 = int(segment_dict['bbox'][0]) + int(segment_dict['bbox'][2])
        y1 = int(segment_dict['bbox'][1]) + int(segment_dict['bbox'][3])
        
        # Scale up to adapt on raw image size
        if ratio != 0:
            x0 = int(x0 // ratio)
            y0 = int(y0 // ratio)
            x1 = int(x1 // ratio)
            y1 = int(y1 // ratio)
        
        # Load mask
        mask = segment_dict['segmentation']
        
        # Process image
        new_image = Image.new('RGB', size=(image_width, image_height), color=(255, 255, 255))
        new_image.paste(Image.fromarray(raw_image), (0, 0), mask=Image.fromarray(mask).resize((image_width, image_height)))
        
        if tag == "mask":
            roi = gen_square_crops(new_image, [x0, y0, x1, y1])  # crop by mask
        elif tag == "bbox":
            roi = gen_square_crops(Image.fromarray(raw_image), [x0, y0, x1, y1])  # crop by bbox
        else:
            raise ValueError("Wrong tag!")
        
        rois.append(roi)
        
        # Save roi and meta data
        os.makedirs(os.path.join(args.output_dir, scene_name), exist_ok=True)
        roi_path = os.path.join(args.output_dir, scene_name, f"{scene_name}_{str(ind).zfill(3)}.png")
        roi.save(roi_path)
        
        # Create roi metadata
        sel_roi = dict()
        sel_roi['roi_id'] = int(ind)
        sel_roi['image_id'] = 0
        sel_roi['bbox'] = [segment_dict['bbox'][0], 
                           segment_dict['bbox'][1], 
                           segment_dict['bbox'][2], 
                           segment_dict['bbox'][3]]
        sel_roi['area'] = np.count_nonzero(mask)
        sel_roi['roi_dir'] = roi_path
        sel_roi['image_dir'] = temp_path
        sel_roi['image_width'] = scene_image.shape[1]
        sel_roi['image_height'] = scene_image.shape[0]
        sel_roi['scale'] = int(1/ratio)
        
        sel_rois.append(sel_roi)
    
    # Save proposals
    with open(os.path.join(args.output_dir, f'proposals_on_{scene_name}.json'), 'w') as f:
        json.dump(sel_rois, f)
    
    # Extract features for proposals
    transform = pth_transforms.Compose([pth_transforms.ToTensor(),])
    scene_dataset = RealWorldDataset(args.output_dir, scene_name, data=rois, transform=transform, imsize=args.imsize)
    scene_features = extract_features(
        scene_dataset, args.batch_size, args.num_workers
    )
    
    # Save scene features
    feat_dict = dict()
    feat_dict['features'] = scene_features.detach().cpu().tolist()
    with open(os.path.join(args.output_dir, f'scene_features_{scene_name}.json'), 'w') as f:
        json.dump(feat_dict, f)
    
    # Normalize features
    scene_features = nn.functional.normalize(scene_features, dim=1, p=2)
    
    # Compute similarity and match proposals
    scene_cnt = [0, *scene_dataset.cfg['length']]
    scene_idx = [sum(scene_cnt[:i + 1]) for i in range(len(scene_cnt))]
    scene_features_list = [scene_features[scene_idx[i]:scene_idx[i + 1]] for i in
                            range(len(scene_dataset.cfg['length']))]
    
    proposals = scene_dataset.cfg['proposals']
    proposals_list = [proposals[scene_idx[i]:scene_idx[i + 1]] for i in range(len(scene_dataset.cfg['length']))]
    
    # 修改这部分来处理不同object有不同数量的example的情况
    num_object = len(object_dataset.cfg['obj_name'])
    
    # 获取每个object对应的example数量
    example_counts = object_dataset.cfg['length']
    
    # 创建索引映射，记录每个object的特征开始和结束位置
    obj_indices = []
    start_idx = 0
    for count in example_counts:
        obj_indices.append((start_idx, start_idx + count))
        start_idx += count
    
    # 对于结果计算部分进行重写
    results = []
    for idx, scene_feature in enumerate(scene_features_list):
        # 获取当前场景的proposals
        proposals = proposals_list[idx]
        
        # 获取object数量和每个object的example数量
        num_object = len(object_dataset.cfg['obj_name'])
        example_counts = object_dataset.cfg['length']
        
        # 创建新的相似度矩阵
        sims = torch.zeros((len(scene_feature), num_object), device=scene_feature.device)
        
        # 跟踪特征的起始索引
        start_idx = 0
        
        # 为每个object计算与所有场景proposals的相似度
        for obj_idx in range(num_object):
            # 获取当前object的example数量
            num_examples = example_counts[obj_idx]
            
            # 获取当前object的所有example特征
            obj_features = object_features[start_idx:start_idx + num_examples]
            
            # 更新起始索引
            start_idx += num_examples
            
            # 计算每个proposal与当前object的所有example的相似度
            # 为每个proposal找到与当前object的最大相似度
            for prop_idx in range(len(scene_feature)):
                prop_feature = scene_feature[prop_idx:prop_idx+1]  # 保持2D tensor
                # 计算当前proposal与当前object的所有example的相似度
                similarities = torch.mm(obj_features, prop_feature.t())  # [num_examples, 1]
                # 取最大相似度
                max_sim, _ = torch.max(similarities, dim=0)
                # 保存到相似度矩阵
                sims[prop_idx, obj_idx] = max_sim.item()
        
        # Stable Matching Strategy
        sel_obj_ids = [str(v) for v in list(np.arange(num_object))]  # ids for selected obj
        sel_roi_ids = [str(v) for v in list(np.arange(len(scene_feature)))]  # ids for selected roi
        
        # Padding
        max_len = max(len(sel_roi_ids), len(sel_obj_ids))
        sel_sims_symmetric = torch.ones((max_len, max_len)) * -1
        sel_sims_symmetric[:len(sel_roi_ids), :len(sel_obj_ids)] = sims.clone()
        
        pad_len = abs(len(sel_roi_ids) - len(sel_obj_ids))
        if len(sel_roi_ids) > len(sel_obj_ids):
            pad_obj_ids = [str(i) for i in range(num_object, num_object + pad_len)]
            sel_obj_ids += pad_obj_ids
        elif len(sel_roi_ids) < len(sel_obj_ids):
            pad_roi_ids = [str(i) for i in range(len(sel_roi_ids), len(sel_roi_ids) + pad_len)]
            sel_roi_ids += pad_roi_ids
        
        # Perform stable matching
        matchedMat = stableMatching(sel_sims_symmetric.detach().data.cpu().numpy())
        predMat_row = np.zeros_like(sel_sims_symmetric.detach().data.cpu().numpy())
        Matches = dict()
        
        for i in range(matchedMat.shape[0]):
            tmp = matchedMat[i, :]
            a = tmp.argmax()
            predMat_row[i, a] = tmp[a]
            Matches[sel_roi_ids[i]] = sel_obj_ids[int(a)]
        
        # Apply threshold
        preds = Matches.copy()
        for key, value in Matches.items():
            # 确保索引在有效范围内
            roi_idx = int(sel_roi_ids.index(key))
            obj_idx = int(sel_obj_ids.index(value))
            
            # 检查索引是否在相似度矩阵范围内
            if roi_idx < sims.shape[0] and obj_idx < sims.shape[1]:
                # 使用原始相似度矩阵进行阈值过滤
                if sims[roi_idx, obj_idx] <= score_threshold:
                    del preds[key]
                    continue
            else:
                # 如果索引超出范围，可能是填充的部分，删除它
                del preds[key]
                continue

        # Save results
        for k, v in preds.items():
            # 确保索引在有效范围内
            if int(k) >= len(proposals) or int(v) >= num_object:
                continue
                
            result = dict()
            result['anns'] = []
            for ann_path in object_dataset.cfg['anns_dir'][int(v)]:
                with open(ann_path, 'r', encoding='utf-8') as f:
                    for line in f:
                        if line.startswith("Object Name:"):
                            result['anns'].append(line.split(":", 1)[1].strip())
            result['anns'] = ' '.join(result['anns'])
            result['image_id'] = proposals[int(k)]['image_id'] 
            result['category_id'] = int(v)
            result['bbox'] = proposals[int(k)]['bbox']
            result['score'] = float(sims[int(k), int(v)])
            result['image_width'] = proposals[int(k)]['image_width']
            result['image_height'] = proposals[int(k)]['image_height']
            result['scale'] = proposals[int(k)]['scale']
            results.append(result)
    
    # Clean up temp file
    try:
        os.unlink(temp_path)
    except:
        pass
    
    # Visualize results
    object_names = object_dataset.cfg['obj_name']

    visualized_img = visualize_detection(np.array(raw_image), results, object_names)
    
    return results, visualized_img, object_names

# ===== FACE DETECTION AND RECOGNITION PART =====


# Initialize face detection and recognition models
@spaces.GPU
def initialize_face_models():
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    mtcnn = MTCNN(
        image_size=160, margin=0, min_face_size=20,
        thresholds=[0.6, 0.7, 0.7], factor=0.709, post_process=True,
        device=device
    )
    resnet = InceptionResnetV1(pretrained='vggface2').eval().to(device)
    return mtcnn, resnet, device

# Get face embeddings from the faces database
def get_face_embeddings(face_dir=FACE_SAVE_PATH):
    mtcnn, resnet, device = initialize_face_models()
    embeddings = []
    face_names = []
    face_paths = []
    face_anns = {}
    
    # Process each face directory in the database
    face_dirs = sorted(glob.glob(os.path.join(face_dir, '*')))
    for face_dir in face_dirs:
        face_paths.append(face_dir)
        same_person_embeddings = []
        face_id = os.path.basename(face_dir)
        
        # Get the first image file for each face
        image_files = sorted(
            glob.glob(os.path.join(face_dir, 'images', '*.png')) +
            glob.glob(os.path.join(face_dir, 'images', '*.jpg')) +
            glob.glob(os.path.join(face_dir, 'images', '*.jpeg')) +
            glob.glob(os.path.join(face_dir, 'images', '*.bmp'))
        )
        if not image_files:
            continue
            
        # Use the first image file to get face name
        ann_files = sorted(glob.glob(os.path.join(face_dir, 'anns', '*.txt')))
        face_name = face_id
        face_names.append(face_name)
        face_anns[face_name] = []
        if ann_files:
            for file in ann_files:
                with open(file, 'r') as f:
                    for line in f:
                        if line.startswith("Object Name:"):
                            face_anns[face_name].append(line.split(":", 1)[1].strip())
        face_anns[face_name] = ' '.join(face_anns[face_name])
        # Process each image for this face
        for img_file in image_files:
            try:
                img = Image.open(img_file).convert('RGB')

                # Convert input image to RGB if needed
                raw_image = np.array(img)
                
                # Detect faces
                boxes, probs = mtcnn.detect(raw_image)


                
                if boxes is not None:
                    # Process each detected face to get embeddings
                    for i, (box, prob) in enumerate(zip(boxes, probs)):
                        if prob < 0.5:  # Minimum confidence for face detection
                            continue
                            
                        # Get coordinates
                        x1, y1, x2, y2 = box.astype(int)
                        
                        # Extract face
                        face = raw_image[y1:y2, x1:x2]
                        img = Image.fromarray(face)


                # Since these are already cropped face images, we might not need MTCNN detection
                # But we'll resize them to the expected size
                img_tensor = pth_transforms.Compose([
                    pth_transforms.Resize((160, 160)),
                    pth_transforms.ToTensor()
                ])(img).unsqueeze(0).to(device)
                
                # Get embedding
                with torch.no_grad():
                    
                    embedding = resnet(img_tensor).detach().cpu().numpy()[0]
                
                same_person_embeddings.append(embedding)
                
                
            except Exception as e:
                print(f"Error processing {img_file}: {str(e)}")
        embeddings.append(np.array(same_person_embeddings))
    return embeddings, face_names, face_paths, face_anns

# Detect and recognize faces in an image
def detect_faces(input_img, score_threshold=0.7):
    mtcnn, resnet, device = initialize_face_models()
    
    # Get reference face embeddings
    face_embeddings, face_names, face_paths, face_anns = get_face_embeddings()
    
    if len(face_embeddings) == 0:
        raw_image = np.array(input_img.convert('RGB'))
        return [], raw_image, []
    
    # Convert input image to RGB if needed
    raw_image = np.array(input_img.convert('RGB'))
    
    # Detect faces
    boxes, probs = mtcnn.detect(raw_image)
  
    results = []
    detected_embeddings = []
    detected_boxes = []
    
    if boxes is not None:
        # Process each detected face to get embeddings
        for i, (box, prob) in enumerate(zip(boxes, probs)):
            if prob < 0.9:  # Minimum confidence for face detection
                continue
                
            # Get coordinates
            x1, y1, x2, y2 = box.astype(int)
            
            try:
                # Extract face
                face = raw_image[y1:y2, x1:x2]
                face_pil = Image.fromarray(face)
                
                # Convert to tensor and get embedding
                face_tensor = pth_transforms.Compose([
                    pth_transforms.Resize((160, 160)),
                    pth_transforms.ToTensor()
                ])(face_pil).unsqueeze(0).to(device)
                
                with torch.no_grad():
                    embedding = resnet(face_tensor).detach().cpu().numpy()[0]
                
                # Store embedding and box for stable matching
                detected_embeddings.append(embedding)
                detected_boxes.append((x1, y1, x2, y2))
            except Exception as e:
                print(f"Error processing face {i}: {str(e)}")
    
        # Use stable matching to find the best matches if we have detected faces
        if detected_embeddings:
            matches, similarities = match_faces_stable_matching(
                face_embeddings, 
                detected_embeddings, 
                score_threshold
            )
            
            # Create results from the matches
            for detected_idx, ref_idx in matches.items():
                x1, y1, x2, y2 = detected_boxes[detected_idx]
                result = {
                    'category_id': ref_idx,
                    'bbox': [x1, y1, x2-x1, y2-y1],
                    'score': float(similarities[detected_idx, ref_idx]),
                    'scale': 1.0,  # Original scale
                    'face_name': face_names[ref_idx],
                    'face_path': face_paths[ref_idx],
                    'face_anns': face_anns[face_names[ref_idx]]
                }
                results.append(result)
    
    # Draw results on image
    visualized_img = visualize_face_detection(raw_image, results, face_names)
    
    return results, visualized_img, face_names


def visualize_face_detection(image, results, face_names):
    """Visualize face detection results on image"""
    output_img = image.copy()
    for res in results:
        x, y, w, h = res['bbox']
        category_id = res['category_id']
        category = res['face_name']
        score = res['score']
        
        # Draw rectangle
        cv2.rectangle(output_img, (x, y), (x+w, y+h), (0, 0, 255), 2)
        
        # Add label
        text = f"{category}: {score:.2f}"
        cv2.putText(output_img, text, (x, y-5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
    
    return output_img

# Function to match faces using stable matching
def match_faces_stable_matching(face_embeddings, detected_embeddings, score_threshold=0.7):
    # 计算相似度矩阵（每个对象的多个示例取最大值）
    similarities = np.zeros((len(detected_embeddings), len(face_embeddings)))
    for i, detect_emb in enumerate(detected_embeddings):
        for j, ref_embs in enumerate(face_embeddings):
            # 对每个对象的多个示例取最大相似度
            max_similarity = 0.
            for ref_emb in ref_embs:
                dist = np.linalg.norm(detect_emb - ref_emb)
                similarity = 1.0 / (dist + 1e-10)  # 余弦相似度近似表示
                if similarity > max_similarity:
                    max_similarity = similarity
            similarities[i, j] = max_similarity

    # 生成对称矩阵并填充虚拟项
    # ------------------------------------------------------------
    sel_obj_ids = [str(j) for j in range(len(face_embeddings))]  # 对象ID列表
    sel_roi_ids = [str(i) for i in range(len(detected_embeddings))]  # ROI ID列表

    max_len = max(len(sel_roi_ids), len(sel_obj_ids))  # 补齐后的长度
    sim_matrix_padded = np.ones((max_len, max_len)) * -1  # 初始化填充矩阵为-1
    sim_matrix_padded[:len(sel_roi_ids), :len(sel_obj_ids)] = similarities  # 填充有效区域

    # 补齐ID列表以匹配矩阵维度
    pad_len = abs(len(sel_roi_ids) - len(sel_obj_ids))
    if len(sel_roi_ids) > len(sel_obj_ids):
        pad_obj_ids = [str(j + len(sel_obj_ids)) for j in range(pad_len)]  # 生成虚拟对象ID
        sel_obj_ids += pad_obj_ids
    elif len(sel_roi_ids) < len(sel_obj_ids):
        pad_roi_ids = [str(i + len(sel_roi_ids)) for i in range(pad_len)]  # 生成虚拟ROI ID
        sel_roi_ids += pad_roi_ids

    # 稳定匹配算法
    # ------------------------------------------------------------
    matched_matrix = stableMatching(sim_matrix_padded)  # 输入补齐后的对称矩阵

    # 解析匹配结果并应用阈值
    # ------------------------------------------------------------
    matches = {}
    for i in range(matched_matrix.shape[0]):
        # 1. 过滤虚拟ROI（超出原始数量的行）
        if i >= len(detected_embeddings):
            continue
        
        # 2. 获取匹配的对象索引
        j = np.argmax(matched_matrix[i, :])
        
        # 3. 过滤虚拟对象（超出原始数量的列）
        if j >= len(face_embeddings):
            continue
        
        # 4. 应用相似度阈值（使用原始未填充的相似度矩阵）
        if similarities[i, j] > score_threshold:
            matches[i] = j  # 保存原始索引的对应关系
    
    return matches, similarities
# 1. Add the combined detection function
def combined_detection(img, obj_threshold, face_threshold, tag):
    """
    Run both object detection and face detection on the same image
    
    Args:
        img: PIL Image to detect objects and faces in
        obj_threshold: Threshold for object detection
        face_threshold: Threshold for face detection
        tag: Proposal type for object detection ("mask" or "bbox")
        
    Returns:
        combined_results: Combined JSON results
        output_img: Image with detection visualizations
    """
    # Run object detection
    obj_results, obj_img, obj_names = detect_objects(img, obj_threshold, tag)
    
    # Run face detection
    face_results, face_img, face_names = detect_faces(img, face_threshold)
    
    # Combine results
    combined_results = {
        "objects": obj_results,
        "faces": face_results
    }
    
    # Create combined visualization
    # We'll use a new image to avoid overlapping visuals from the two separate functions
    raw_image = np.array(img.convert('RGB'))
    combined_img = raw_image.copy()
    i = 1
    # Draw object detections (green boxes)
    for res in obj_results:
        x, y, w, h = res['bbox']
        category = obj_names[res['category_id']]
        score = res['score']
        
        # Convert to absolute coordinates based on scale
        x = int(x * res['scale'])
        y = int(y * res['scale'])
        w = int(w * res['scale'])
        h = int(h * res['scale'])
        
        # Draw rectangle
        cv2.rectangle(combined_img, (x, y), (x+w, y+h), (0, 0, 255), 1)
        
        # Add label
        text = f"[{i}]: {score:.2f}"
        i = i+1
        # 创建一个覆盖层
        overlay = combined_img.copy()
        (text_width, text_height), _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
        cv2.rectangle(overlay, (x, y-text_height-5), (x + text_width, y+5), (255, 255, 255), -1)

        # 合并覆盖层与原图（调整透明度）
        alpha = 0.7  # 透明度参数：0 完全透明，1 完全不透明
        cv2.addWeighted(overlay, alpha, combined_img, 1-alpha, 0, combined_img)

        # 绘制文字
        cv2.putText(combined_img, text, (x, y-5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1, cv2.LINE_AA)
    
    # Draw face detections (red boxes)
    for res in face_results:
        x, y, w, h = res['bbox']
        category = res['face_name']
        score = res['score']
        
        # Draw rectangle
        cv2.rectangle(combined_img, (x, y), (x+w, y+h), (0, 0, 255), 1)
        
        # Add label
        text = f"[{i}]: {score:.2f}"
        i = i+1
        #cv2.putText(combined_img, text, (x, y-5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1, cv2.LINE_AA)
        # 创建一个覆盖层
        overlay = combined_img.copy()
        (text_width, text_height), _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
        cv2.rectangle(overlay, (x, y-text_height-5), (x + text_width, y+5), (255, 255, 255), -1)

        # 合并覆盖层与原图（调整透明度）
        alpha = 0.7  # 透明度参数：0 完全透明，1 完全不透明
        cv2.addWeighted(overlay, alpha, combined_img, 1-alpha, 0, combined_img)

        # 绘制文字
        cv2.putText(combined_img, text, (x, y-5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1, cv2.LINE_AA)
    return combined_results, combined_img


################################################
# Import necessary libraries
import torch
from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
from qwen_vl_utils import process_vision_info

# Load model and processor at the application level for reuse
def load_qwen2vl_model():
    model = Qwen2VLForConditionalGeneration.from_pretrained(
        "weihongliang/RC-Qwen2VL-7b",
        torch_dtype=torch.bfloat16,
        device_map="cuda"
    )
    min_pixels = 256 * 28 * 28
    max_pixels = 1280 * 28 * 28
    processor = AutoProcessor.from_pretrained(
        "Qwen/Qwen2-VL-2B-Instruct", min_pixels=min_pixels, max_pixels=max_pixels
    )
    #processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct")
    return model, processor

# Try to load the model, but handle errors if it fails
try:
    qwen_model, qwen_processor = load_qwen2vl_model()
    qwen_model_loaded = True
except Exception as e:
    print(f"Failed to load Qwen2-VL model: {e}")
    qwen_model_loaded = False

# Function to process detection results and use Qwen2-VL for answering questions
@spaces.GPU

def ask_qwen_about_detections(input_image, question, obj_threshold, face_threshold, tag):
    """
    Process an image with detection and use Qwen2-VL to answer questions
    """
    # Check if the model is loaded
    if not qwen_model_loaded:
        return "Qwen2-VL model not loaded. Please check console for errors.", None, None
    
    # Get detection results and formatted text
    qwen_input, output_img = process_image_for_qwen(input_image, obj_threshold, face_threshold, tag)
    print(qwen_input)
    print(input_image.size)
    input_image.save('./temp_image.jpg')
    # Prepare input for Qwen2-VL
    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "image",
                    "image": './temp_image.jpg',
                },
                {
                    "type": "text", 
                    "text": f"{qwen_input}\nAnswer the following question based on the information above and the given image, and provide citations for your response.\n{question}"
                },
            ],
        }
    ]
    
    # Apply chat template
    text = qwen_processor.apply_chat_template(
        messages, 
        tokenize=False, 
        add_generation_prompt=True
    )
    
    # Process vision info
    image_inputs, video_inputs = process_vision_info(messages)
    print(image_inputs)
    # Prepare inputs
    inputs = qwen_processor(
        text=[text], 
        images=image_inputs, 
        videos=None, 
        padding=True, 
        return_tensors="pt"
    )
    
    # Move to GPU if available
    device = "cuda" if torch.cuda.is_available() else "cpu"
    inputs = inputs.to(device)
    
    # Generate answer
    with torch.no_grad():
        generated_ids = qwen_model.generate(**inputs, max_new_tokens=10000)
    
    generated_ids_trimmed = [
        out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
    ]
    
    # Decode the answer
    answer = qwen_processor.batch_decode(
        generated_ids_trimmed, 
        skip_special_tokens=True, 
        clean_up_tokenization_spaces=False
    )[0]
    
    return answer, output_img, qwen_input

def format_for_qwen(results, image_width, image_height):
    """
    Format the detection results for Qwen2-VL model input
    
    Args:
        results: Combined detection results from combined_detection function
        image_width: Width of the original image
        image_height: Height of the original image
        
    Returns:
        formatted_text: Text formatted in the required pattern for Qwen2-VL
    """
    # Combine object and face detections
    all_detections = []
    print(results)
    # Add object detections
    for i, obj in enumerate(results["objects"]):
        bbox = obj["bbox"]
        # Convert bbox [x, y, w, h] to [x1, y1, x2, y2] format
        x1, y1 = bbox[0], bbox[1]
        x2, y2 = x1 + bbox[2], y1 + bbox[3]
        
        # Apply scaling if present in the object
        scale = obj.get("scale", 1.0)
        x1 = int(x1 * scale)
        y1 = int(y1 * scale)
        x2 = int(x2 * scale)
        y2 = int(y2 * scale)
        
        # Normalize coordinates and multiply by 1000 as required
        norm_x1 = int((x1 / image_width) * 1000)
        norm_y1 = int((y1 / image_height) * 1000)
        norm_x2 = int((x2 / image_width) * 1000)
        norm_y2 = int((y2 / image_height) * 1000)
        
        #category = results["objects"][i]["category_id"]
        category_name = 'object'
        
        all_detections.append({
            "box": [(norm_x1, norm_y1), (norm_x2, norm_y2)],
            "info": remove_numbers_in_brackets(obj['anns']),
            "type": 'object'
        })
    
    # Add face detections
    for i, face in enumerate(results["faces"]):
        bbox = face["bbox"]
        # Face bbox is already in [x, y, w, h] format
        x1, y1 = bbox[0], bbox[1]
        x2, y2 = x1 + bbox[2], y1 + bbox[3]
        
        # Normalize coordinates and multiply by 1000 as required
        norm_x1 = int((x1 / image_width) * 1000)
        norm_y1 = int((y1 / image_height) * 1000)
        norm_x2 = int((x2 / image_width) * 1000)
        norm_y2 = int((y2 / image_height) * 1000)
        
        all_detections.append({
            "box": [(norm_x1, norm_y1), (norm_x2, norm_y2)],
            "info": remove_numbers_in_brackets(face['face_anns']),
            "type": "person"
        })
    
    # Format the detection results as required
    formatted_text = ""
    for i, detection in enumerate(all_detections):
        box = detection["box"]
        info = detection["info"]
        formatted_text += f"[{i+1}]: The information of the {detection['type']} located at <|box_start|>({box[0][0]},{box[0][1]}),({box[1][0]},{box[1][1]})<|box_end|> in the image: {info}\n"
    
    return formatted_text

# Example of how to integrate this with your combined_detection function
def process_image_for_qwen(img, obj_threshold, face_threshold, tag):
    """
    Process an image through detection and format results for Qwen2-VL
    
    Args:
        img: PIL Image to process
        obj_threshold: Threshold for object detection
        face_threshold: Threshold for face detection
        tag: Proposal type for object detection ("mask" or "bbox")
        
    Returns:
        qwen_input: Formatted text for Qwen2-VL model
        output_img: Image with detection visualizations
    """
    # Get image dimensions
    width, height = img.size
    
    # Run combined detection
    results, output_img = combined_detection(img, obj_threshold, face_threshold, tag)
    
    # Format results for Qwen2-VL
    qwen_input = format_for_qwen(results, width, height)
    
    return qwen_input, output_img
# Add this function to display detection info without asking a question
def show_detection_info(input_image, obj_threshold, face_threshold, tag):
    """
    Process the image through detection and show the formatted results
    """
    if input_image is None:
        return "Please upload an image first", None
    
    # Process the image to get detection info
    qwen_input, output_img = process_image_for_qwen(input_image, obj_threshold, face_threshold, tag)
    
    return qwen_input, output_img
# Add this function to your Gradio interface
def detect_and_format_for_qwen(input_image, obj_threshold, face_threshold, tag):
    """
    Detect objects/faces and format results for Qwen2-VL
    """
    qwen_input, output_img = process_image_for_qwen(input_image, obj_threshold, face_threshold, tag)
    return qwen_input, output_img
# Create a directory for example images if it doesn't exist
EXAMPLE_DIR = "./examples"
os.makedirs(EXAMPLE_DIR, exist_ok=True)

# Function to create a simple placeholder image if needed
def create_placeholder_image(filename, width=500, height=500, color=(100, 150, 200)):
    """Create a simple colored image as a placeholder"""
    img = Image.new('RGB', (width, height), color=color)
    img.save(filename)
    return filename

# Example person and object images - in real deployment, replace these with actual example images
def ensure_example_images():
    """Ensure example images exist, creating placeholders if needed"""
    examples = {
        "person1.jpg": (500, 600, (220, 180, 170)),
        "person2.jpg": (500, 600, (200, 170, 160)),
        "object1.jpg": (600, 400, (180, 200, 220)),
        "object2.jpg": (600, 400, (160, 220, 190)),
        "scene1.jpg": (800, 600, (170, 190, 210)),
        "scene2.jpg": (800, 600, (190, 210, 180))
    }
    
    image_paths = {}
    for name, (width, height, color) in examples.items():
        path = os.path.join(EXAMPLE_DIR, name)
        if not os.path.exists(path):
            create_placeholder_image(path, width, height, color)
        image_paths[name] = path
    
    return image_paths

# Prepare example images
example_image_paths = ensure_example_images()

# Example data for the first tab (Upload Multimodal Personalized Information)
tab1_examples = [
    # Person examples
    [
        "👤face",                                # Mode selection
        "./examples/hrx.jpeg",   
        ["./examples/hrx.jpeg"],  # Image input
        """Jen-Hsun "Jensen" Huang (Chinese: 黃仁勳; pinyin: Huáng Rénxūn; Pe̍h-ōe-jī: N̂g Jîn-hun; born February 17, 1963) is a Taiwanese and American businessman, electrical engineer, and philanthropist who is the president, co-founder, and chief executive officer (CEO) of Nvidia, the world's largest semiconductor company. In February 2025, Forbes estimated Huang's net worth at US$114.5 billion, making him the 14th wealthiest person in the world."""  # Personal info
    ],

    
    # Object examples

    [
        "📦object",                              # Mode selection
        "./examples/3080.jpeg",  # Image input
        ["./examples/3080.jpeg"],
        "The GeForce RTX™ 3080 delivers the ultra performance that gamers crave, powered by Ampere—NVIDIA’s 2nd gen RTX architecture. It’s built with enhanced RT Cores and Tensor Cores, new streaming multiprocessors, and superfast G6X memory for an amazing gaming experience."  # Object info
    ],
    [
        "👤face",                                # Mode selection
        "./musk.jpeg",  
        ["./musk.jpeg"],  # Image input
        "Elon Reeve Musk (/ˈiːlɒn/ EE-lon; born June 28, 1971) is a businessman known for his leadership of Tesla, SpaceX, and X (formerly Twitter). Since 2025, he has been a senior advisor to United States President Donald Trump and the de facto head of the Department of Government Efficiency (DOGE). Musk is the wealthiest person in the world; as of March 2025, Forbes estimates his net worth to be US$345 billion. He was named Time magazine's Person of the Year in 2021."
    ],
    [
        "📦object",                              # Mode selection
        "./cybertruck.jpg",  # Image input
        ["./cybertruck.jpg"],
        "The Tesla Cybertruck is a battery-powered electric pickup truck manufactured by Tesla, Inc. since 2023.[6] Introduced as a concept vehicle in November 2019, its body design is reminiscent of low-polygon modeling, consisting of flat stainless steel sheet panels."
    ]
]

# Example data for the second tab (Personalized Multimodal Understanding)
tab2_examples = [
    [
        "./examples/hrx_3080.jpg",                           # Image input
        "Who is in this image and what are they doing?"              # Question
    ],
    [
        "./musk_cybertruck.jpeg",                           # Image input
        "Describe the image."              # Question
    ],
    [
        "./musk_and_huang.jpg",                           # Image input
        "Describe the image."              # Question
    ],
]

# Function to clear the directory when a new image is uploaded
def clear_directory(image):
    directory_path = "./exps"
    try:
        # Check if directory exists
        if os.path.exists(directory_path):
            # Remove all files and subdirectories
            for item in os.listdir(directory_path):
                item_path = os.path.join(directory_path, item)
                if os.path.isfile(item_path):
                    os.remove(item_path)
                elif os.path.isdir(item_path):
                    shutil.rmtree(item_path)
            print(f"Cleared directory: {directory_path}")
        else:
            print(f"Directory does not exist: {directory_path}")
    except Exception as e:
        print(f"Error clearing directory: {e}")
    
    # Return the image to be used in the next function in the chain if needed
    return image

# Function to handle clicks on examples in Tab 1 (REVISED SIGNATURE)
def handle_example_click(mode, img_path_display, file_list_for_state, obj_info):
    """
    Handles the click event on a gr.Examples row for Tab 1.
    Updates the global state and returns values to populate UI components.
    
    Args:
        mode: Value from the first column of the example.
        img_path_display: Value from the second column (often used for the Image component).
        file_list_for_state: Value from the third column (used to update state.images).
        obj_info: Value from the fourth column.
    
    Returns:
        A tuple/list of values for the components specified in gr.Examples outputs.
    """
    print('魏洪亮')
    global state # Ensure we modify the global state object
    
    # --- No longer need to unpack example_data ---
    # if not isinstance(example_data, list) or len(example_data) != 4: ...

    # 1. Update state mode
    state.mode = mode
    print(f"Example click: Mode set to {state.mode}")
    
    # 2. Update state images using the file list
    #    state.add_images resets the index to 0 internally
    try:
        # Ensure file_list_for_state is actually a list of paths/objects
        if isinstance(file_list_for_state, str):
             file_list_for_state = [file_list_for_state]
        elif not isinstance(file_list_for_state, list):
             print(f"Warning: Expected list for file_list_for_state, got {type(file_list_for_state)}. Attempting to use img_path_display.")
             # Use img_path_display as a fallback if file_list isn't a list
             if isinstance(img_path_display, str):
                 file_list_for_state = [img_path_display]
             else:
                 # If img_path_display isn't a string path either, we have a problem
                 print(f"Error: Cannot determine image list for state update.")
                 file_list_for_state = [] # Set to empty list to avoid further errors

        count = state.add_images(file_list_for_state)
        print(f"Example click: Added {count} images to state. Current index: {state.current_image_index}")

    except Exception as e:
         print(f"Error processing example images in state: {e}")
         count = 0

    if count == 0:
        print("Example click: No valid images loaded into state.")
        # Return default/empty values for outputs matching the 'outputs' list length
        return mode, None, None, obj_info, None, "Error loading example image(s).", []

    # 3. Get the current image (which should be the first one after add_images)
    current_image_to_display = state.get_current_image()
    if current_image_to_display is None:
         print("Error: State has images, but get_current_image returned None.")
         return mode, None, None, obj_info, None, "Internal error getting example image.", state.get_gallery()
         
    print(f"Example click: Current image name from state: {state.get_current_image_name()}")

    # 4. Update status text based on the now valid state
    status = state.get_status_text()
    print(f"Example click: Status text: {status}")

    # 5. Determine the object name placeholder based on mode
    object_name_placeholder = "This person is" if mode == "👤face" else "This object is"
    final_object_info = obj_info if obj_info and obj_info.strip() else object_name_placeholder

    # 6. Return values for all output components defined in gr.Examples outputs
    #    Order MUST match the outputs list:
    #    [mode_selection, input_image, file_output, object_name, masked_image, status_text, gallery]
    return (
        mode,                       # For mode_selection
        current_image_to_display,   # For input_image (state ensures this is np array)
        None,                       # For file_output (clear it)
        final_object_info,          # For object_name
        None,                       # For masked_image (clear it)
        status,                     # For status_text
        state.get_gallery()         # For gallery
    )
# Modified app definition to include examples
with gr.Blocks() as app:
    #gr.Markdown("# Personalized Multimodal Understanding with RC-MLLM")
    gr.Markdown("<div style='text-align: center;'><h1 style=' font-size: 28px; '>Personalized Multimodal Understanding with RC-MLLM</h1></div>")
    gr.Markdown("**RC-MLLM** model is developed based on the Qwen2-VL model through a novel method called **RCVIT (Region-level Context-aware Visual Instruction Tuning)**, using the specially constructed **RCMU dataset** for training. Its core feature is the capability for **Region-level Context-aware Multimodal Understanding (RCMU)**. This means it can simultaneously understand both the visual content of specific regions/objects within an image and their associated textual information (utilizing bounding boxes coordinates), allowing it to respond to user instructions in a more context-aware manner. Simply put, RC-MLLM not only understands images but can also integrate the textual information linked to specific objects within the image for understanding. It achieves outstanding performance on RCMU tasks and is suitable for applications like personalized conversation.")

    markdown_content = """
        📑 [Region-Level Context-Aware Multimodal Understanding](https://arxiv.org/abs/2508.12263) | 
        🤗 Models:[RC-Qwen2VL-2b](https://huggingface.co/weihongliang/RC-Qwen2VL-2b/blob/main/README.md) [RC-Qwen2VL-7b](https://huggingface.co/weihongliang/RC-Qwen2VL-7b/blob/main/README.md)| 
        📁 [Dataset](https://huggingface.co/your-model-name) |
        [Github](https://github.com/hongliang-wei/RC-MLLM) | 
        🚀 [Celebrity Recognition and VQA Demo](https://huggingface.co/spaces/weihongliang/RCMLLM)
        """
    gr.Markdown(markdown_content)

    #gr.Markdown("[![arXiv](https://img.shields.io/badge/📄_arXiv-B31B1B.svg?&style=flat-square&logo=arXiv&logoColor=white)](https://arxiv.org/abs/2301.00000)[![arXiv](https://img.shields.io/badge/📄_arXiv-B31B1B.svg?&style=flat-square&logo=arXiv&logoColor=white)](https://arxiv.org/abs/2301.00000)[![GitHub](https://img.shields.io/badge/💻_GitHub-100000?style=flat-square&logo=github&logoColor=white)](https://githucom/username/repository)[![Hugging Face Models](https://img.shields.io/badge/🤗_Models-FF9D00?style=flat-square&logo=huggingface&logoColor=white(https://huggingface.co/models/username/model-name)[![Hugging Face Space](https://img.shields.io/badge/🚀_Space-7358FF?style=flat-square&logo=huggingface&logoColor=white(https://huggingface.co/spaces/username/space-name)")
    gr.Markdown("### 📌 First build a multimodal personalized knowledge base, then perform personalized multimodal understanding with RC-MLLM")

    # Section 1: Upload Multimodal Personalized Information (formerly Tab 1)
    #gr.Markdown("<h2 style='color: #28a745; background-color: #F8F9FA; padding: 10px; border-left: 5px solid #28a745; border-radius: 5px;'>1. Upload Multimodal Personalized Information")

    #gr.Markdown("📖 Upload images, click on people or objects in the images and fill in their personalized information, then save them to create a multimodal personalized knowledge base")
    gr.Markdown("<h2 style='color: #28a745; background-color: #F8F9FA; padding: 10px; border-left: 5px solid #28a745; border-radius: 5px;'>1. Build Multimodal Personalized Knowledge Base<br><span style='font-size: 0.9rem;'>📖 Upload images, click on people or objects in the images and fill in their personalized information, then save them to create a multimodal personalized knowledge base</span></h2>")
    
    # First Row: Upload controls on left, personalized info and save button on right
    with gr.Row():
        # Left column for upload and navigation
        with gr.Column(scale=1):
            mode_selection = gr.Radio(
                ["📦object", "👤face"], 
                label="Object Image or Face Image (Select the type of image to upload)", 
                value="📦object",
            )
            file_output = gr.File(label="Upload Images", file_count="multiple")
            print(file_output)
            
            # Replace buttons with radio selection
            navigation_selection = gr.Radio(
                ["Different Instance", "Same Instance"], 
                label="Support multiple images per instance.", 
                value="Different Instance"
            )
            #next_image_button = gr.Button("Next Image")
        
        # Middle column for images and mask
        with gr.Column(scale=2):
            # Add a markdown component with the instruction text
            gr.Markdown("<span style='color: red; font-weight: bold;'>Click on people or objects in the image to get a mask</span>")
            
            # Images section
            with gr.Row():
                input_image = gr.Image(
                    label="Current Image", 
                    interactive=False
                )
                masked_image = gr.Image(label="Mask")
        
        # Right column for personalized info, status and save button
        with gr.Column(scale=1):
            # Add mode selection radio button at the top of the right column
            """mode_selection = gr.Radio(
                ["📦object", "👤face"], 
                label="Object Image or Face Image (Select the type of image to upload)", 
                value="📦object",
            )"""
            
            object_name = gr.Textbox(label="Input Personalized Information", 
                                placeholder="Enter personalized information of person/object",
                                value="This object is")
            
            status_text = gr.Textbox(label="Status", value="No images loaded")
            
            save_button = gr.Button("Save Multimodal Personalized Information")
    
    # Second Row: Gallery with size limitation
    with gr.Row():
        gallery = gr.Gallery(
            label="Multimodal Personalized Knowledge Base", 
            show_label=True,
            object_fit="contain",
            height="300px",  # Fixed height
            columns=4,       # Set number of columns
            rows=2,          # Set number of rows
            preview=True,    # Show a larger preview on click
            elem_id="limited_gallery"
        )
    
    # Examples section
    tab1_examples_component = gr.Examples(
        examples=tab1_examples,
        inputs=[mode_selection, input_image, file_output, object_name],
        outputs=[
            mode_selection,
            input_image,
            file_output,
            object_name,
            masked_image,
            status_text,
            gallery
        ],
        fn=handle_example_click,
        run_on_click=True,
        label="Examples for information upload",
        #cache_examples=False  # Disable caching to prevent JSON decode errors

    )
    
    # Instructions
    #gr.Markdown("### Instructions: 1. Select mode (object or face) 2. Upload images 3. Click on an item to generate mask 4. Enter item name and save 5. Choose 'Same Instance' or 'Different Instance' for next image 6. Click 'Next Image' to proceed 7. All processed items appear in gallery 8. Objects and faces saved in separate directories")

    # Add a separator between sections
    gr.Markdown("---")
    
    # Section 2: RC-MLLM Integration (formerly Tab 2)
    #gr.Markdown("## 2. Personalized Multimodal Understanding with RC-MLLM")
    #gr.Markdown("<h2 style='color: #28a745; background-color: #F8F9FA; padding: 10px; border-left: 5px solid #28a745; border-radius: 5px;'>2. Personalized Multimodal Understanding with RC-MLLM</h2>")
    #gr.Markdown("<h2 style='border-bottom: 3px solid #28a745; color: #28a745; font-weight: bold; padding-bottom: 5px;'>2. Personalized Multimodal Understanding with RC-MLLM</h2>")
    #gr.Markdown("📌 Upload images and use the RC-MLLM model for personalized Q&A")
    gr.Markdown("<h2 style='color: #28a745; background-color: #F8F9FA; padding: 10px; border-left: 5px solid #28a745; border-radius: 5px;'>2. Personalized Multimodal Understanding with RC-MLLM<br><span style='font-size: 0.9rem;'>📖 Upload images and use the RC-MLLM model for personalized Q&A</span></h2>")
    with gr.Row():
        with gr.Column():
            qwen_input_image = gr.Image(type="pil", label="Input Image")
            # Set up the change event to trigger the directory clearing
            qwen_input_image.change(
                fn=clear_directory,
                inputs=[qwen_input_image],
                #outputs=[qwen_input_image]  # Pass through the image
            )
            with gr.Row():
                with gr.Column():
                    qwen_obj_threshold = gr.Slider(
                        minimum=0.0, maximum=1.0, value=0.6, step=0.01, 
                        label="Object Detection Threshold"
                    )
                    qwen_tag_choice = gr.Radio(
                        choices=["mask", "bbox"], value="mask", 
                        label="Object Proposal Type",
                        visible=False
                    )
                
                with gr.Column():
                    qwen_face_threshold = gr.Slider(
                        minimum=0.0, maximum=1.5, value=0.7, step=0.01, 
                        label="Face Detection Threshold"
                    )
            
            qwen_question = gr.Textbox(
                label="Question", 
                placeholder="Ask a question about the objects/faces in the image...",
                lines=2
            )
            
            qwen_ask_button = gr.Button("Ask RC-MLLM-7B")
        
        with gr.Column():
            qwen_output_image = gr.Image(label="Detection Result")
            
            # Add new textarea to display the formatted detection information
            qwen_input_display = gr.Textbox(
                label="Detection Information", 
                lines=6,
                max_lines=15,
                interactive=False
            )
            
            qwen_answer = gr.Textbox(
                label="RC-MLLM Answer", 
                lines=8,
                max_lines=15
            )
    
    # Add Examples for MLLM Section
    gr.Examples(
        examples=tab2_examples,
        inputs=[qwen_input_image, qwen_question],
        label="Examples for visual question answering",
        #cache_examples=False  # Disable caching to prevent JSON decode errors

    )
    
    # Model status display
    model_status = gr.Markdown(
        "✅ RC-MLLM model loaded successfully" if qwen_model_loaded else 
        "❌ RC-MLLM model not loaded. Please check console for errors."
    )
    
    # Instructions for RC-MLLM section
    #gr.Markdown("### Instructions: 1. Upload an image 2. Adjust detection thresholds 3. Enter a question 4. Click 'Ask RC-MLLM-7B' 5. View analysis results")
    
    # Event handler for the MLLM question answering
    qwen_ask_button.click(
        fn=ask_qwen_about_detections,
        inputs=[
            qwen_input_image,
            qwen_question, 
            qwen_obj_threshold, 
            qwen_face_threshold, 
            qwen_tag_choice
        ],
        outputs=[qwen_answer, qwen_output_image, qwen_input_display]
    )
    
    # Event handlers for Section 1 (Segmentation)
    # Add handler for mode selection
    mode_selection.change(update_mode, inputs=[mode_selection], outputs=[status_text])
    
    # Add JavaScript to update the textbox value based on mode selection
    mode_selection.change(
        fn=lambda mode: "This person is" if mode == "👤face" else "This object is",
        inputs=mode_selection,
        outputs=object_name
    )
    
    # Modified upload_images function wrapper that also updates the object_name based on mode
    def upload_and_set_name(file_output, mode):
        # First call the original upload_images function
        input_image, masked_image, status_text, gallery = upload_images(file_output)
        # Then set the object_name based on the current mode
        object_name_value = "This person is" if mode == "👤face" else "This object is"
        return input_image, masked_image, status_text, gallery, object_name_value
    
    # Modified file_output.upload event handler
    file_output.upload(
        upload_and_set_name, 
        inputs=[file_output, mode_selection], 
        outputs=[input_image, masked_image, status_text, gallery, object_name]
    )
    
    # New combined navigation function that uses the radio selection
    def navigate_with_selection(navigation_type):
        is_same_object = (navigation_type == "Same Instance")
        return navigate_images(is_same_object)
    
    # Event handler for the next image button
    """next_image_button.click(
        navigate_with_selection, 
        inputs=[navigation_selection], 
        outputs=[input_image, masked_image, status_text, gallery, file_output]
    )"""
    
    input_image.select(
        generate_mask,
        inputs=[input_image],
        outputs=[masked_image, status_text, gallery]
    )
    
    save_button.click(save_mask_and_text, inputs=[object_name], outputs=[status_text, gallery])

# Run the app
if __name__ == "__main__":
    app.launch()