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("

Personalized Multimodal Understanding with RC-MLLM

") 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("

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("

1. Build Multimodal Personalized Knowledge Base
📖 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

") # 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("Click on people or objects in the image to get a mask") # 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("

2. Personalized Multimodal Understanding with RC-MLLM

") #gr.Markdown("

2. Personalized Multimodal Understanding with RC-MLLM

") #gr.Markdown("📌 Upload images and use the RC-MLLM model for personalized Q&A") gr.Markdown("

2. Personalized Multimodal Understanding with RC-MLLM
📖 Upload images and use the RC-MLLM model for personalized Q&A

") 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()