import os from pathlib import Path import gradio as gr import numpy as np import spaces import torch from detection_viewer import DetectionViewer from PIL import Image from transformers import ( AutoImageProcessor, AutoModelForZeroShotObjectDetection, AutoProcessor, Mask2FormerForUniversalSegmentation, RTDetrForObjectDetection, RTDetrImageProcessor, VitPoseForPoseEstimation, ) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") USE_SMALL_MODELS = os.environ.get("USE_SMALL_MODELS", "0") == "1" ASSETS_DIR = Path(__file__).parent / "assets" # ============================================================ # RT-DETR (Object Detection) # ============================================================ RTDETR_MODEL_ID = "PekingU/rtdetr_r18vd" if USE_SMALL_MODELS else "PekingU/rtdetr_r101vd" rtdetr_processor = RTDetrImageProcessor.from_pretrained(RTDETR_MODEL_ID) rtdetr_model = RTDetrForObjectDetection.from_pretrained(RTDETR_MODEL_ID).eval().to(device) # ============================================================ # Grounding DINO (Zero-Shot Detection) # ============================================================ GDINO_MODEL_ID = "IDEA-Research/grounding-dino-tiny" if USE_SMALL_MODELS else "IDEA-Research/grounding-dino-base" gdino_processor = AutoProcessor.from_pretrained(GDINO_MODEL_ID) gdino_model = AutoModelForZeroShotObjectDetection.from_pretrained(GDINO_MODEL_ID).eval().to(device) # ============================================================ # Mask2Former (Instance Segmentation) # ============================================================ M2F_MODEL_ID = ( "facebook/mask2former-swin-tiny-coco-instance" if USE_SMALL_MODELS else "facebook/mask2former-swin-large-coco-instance" ) m2f_processor = AutoImageProcessor.from_pretrained(M2F_MODEL_ID) m2f_model = Mask2FormerForUniversalSegmentation.from_pretrained(M2F_MODEL_ID).eval().to(device) # ============================================================ # ViTPose (Pose Estimation) # ============================================================ VITPOSE_DET_MODEL_ID = "PekingU/rtdetr_r18vd" if USE_SMALL_MODELS else "PekingU/rtdetr_r50vd" vitpose_det_processor = AutoProcessor.from_pretrained(VITPOSE_DET_MODEL_ID) vitpose_det_model = RTDetrForObjectDetection.from_pretrained(VITPOSE_DET_MODEL_ID).eval().to(device) PERSON_LABEL_ID = next(k for k, v in vitpose_det_model.config.id2label.items() if v == "person") VITPOSE_MODEL_ID = "usyd-community/vitpose-base-simple" if USE_SMALL_MODELS else "usyd-community/vitpose-plus-large" VITPOSE_IS_MOE = not USE_SMALL_MODELS # vitpose-plus models use Mixture-of-Experts heads vitpose_processor = AutoProcessor.from_pretrained(VITPOSE_MODEL_ID) vitpose_model = VitPoseForPoseEstimation.from_pretrained(VITPOSE_MODEL_ID).eval().to(device) SKELETON = vitpose_model.config.edges COCO_KEYPOINT_NAMES = [ "nose", "left_eye", "right_eye", "left_ear", "right_ear", "left_shoulder", "right_shoulder", "left_elbow", "right_elbow", "left_wrist", "right_wrist", "left_hip", "right_hip", "left_knee", "right_knee", "left_ankle", "right_ankle", ] # ============================================================ # Utilities # ============================================================ def _mask_to_rle(mask: np.ndarray) -> dict: """Convert a binary mask to uncompressed RLE (column-major / COCO format).""" h, w = mask.shape flat = mask.ravel(order="F") changes = np.diff(flat) change_idx = np.flatnonzero(changes) runs = np.diff(np.concatenate([[-1], change_idx, [len(flat) - 1]])) counts = runs.tolist() if flat[0] == 1: counts = [0, *counts] return {"counts": counts, "size": [h, w]} def _mask_to_bbox(mask: np.ndarray) -> dict: """Compute bounding box from a binary mask.""" ys, xs = np.where(mask) x_min, x_max = int(xs.min()), int(xs.max()) y_min, y_max = int(ys.min()), int(ys.max()) return {"x": x_min, "y": y_min, "width": x_max - x_min + 1, "height": y_max - y_min + 1} # ============================================================ # Inference helpers # ============================================================ def _detect_rtdetr(image: Image.Image, _labels: str, threshold: float) -> tuple[Image.Image, list[dict], dict]: inputs = rtdetr_processor(images=image, return_tensors="pt").to(device) outputs = rtdetr_model(**inputs) results = rtdetr_processor.post_process_object_detection( outputs, target_sizes=torch.tensor([(image.height, image.width)]), threshold=threshold, )[0] annotations = [] for score, label_id, box in zip(results["scores"], results["labels"], results["boxes"], strict=True): x_min, y_min, x_max, y_max = box.tolist() annotations.append( { "bbox": {"x": x_min, "y": y_min, "width": x_max - x_min, "height": y_max - y_min}, "score": round(score.item(), 3), "label": rtdetr_model.config.id2label[label_id.item()], } ) return image, annotations, {"score_threshold": (threshold, 1.0)} def _detect_gdino(image: Image.Image, labels: str, threshold: float) -> tuple[Image.Image, list[dict], dict]: text = labels.strip().rstrip(".") candidate_labels = [part.strip() for part in text.split(",") if part.strip()] text_prompt = ". ".join(candidate_labels) + "." inputs = gdino_processor(images=image, text=text_prompt, return_tensors="pt").to(device) outputs = gdino_model(**inputs) results = gdino_processor.post_process_grounded_object_detection( outputs, input_ids=inputs["input_ids"], target_sizes=[(image.height, image.width)], threshold=threshold, text_threshold=threshold, )[0] annotations = [] for score, label, box in zip(results["scores"], results["labels"], results["boxes"], strict=True): x_min, y_min, x_max, y_max = box.tolist() annotations.append( { "bbox": {"x": x_min, "y": y_min, "width": x_max - x_min, "height": y_max - y_min}, "score": round(float(score), 3), "label": label, } ) return image, annotations, {"score_threshold": (threshold, 1.0)} def _detect_m2f(image: Image.Image, _labels: str, threshold: float) -> tuple[Image.Image, list[dict], dict]: inputs = m2f_processor(images=image, return_tensors="pt").to(device) outputs = m2f_model(**inputs) results = m2f_processor.post_process_instance_segmentation( outputs, target_sizes=[(image.height, image.width)], threshold=threshold, )[0] segmentation = results["segmentation"].cpu().numpy().astype(np.uint8) annotations = [] for segment in results["segments_info"]: binary_mask = (segmentation == segment["id"]).astype(np.uint8) if binary_mask.sum() == 0: continue annotations.append( { "bbox": _mask_to_bbox(binary_mask), "mask": _mask_to_rle(binary_mask), "score": round(float(segment["score"]), 3), "label": m2f_model.config.id2label[int(segment["label_id"])], } ) return image, annotations, {"score_threshold": (threshold, 1.0)} def _detect_vitpose( image: Image.Image, _labels: str, threshold: float ) -> tuple[Image.Image, list[dict], dict] | tuple[Image.Image, list]: # Step 1: Detect persons with RT-DETR det_inputs = vitpose_det_processor(images=image, return_tensors="pt").to(device) det_outputs = vitpose_det_model(**det_inputs) det_results = vitpose_det_processor.post_process_object_detection( det_outputs, target_sizes=torch.tensor([(image.height, image.width)]), threshold=threshold, )[0] person_mask = det_results["labels"] == PERSON_LABEL_ID person_boxes_voc = det_results["boxes"][person_mask].cpu().numpy() person_scores = det_results["scores"][person_mask].cpu().numpy() if len(person_boxes_voc) == 0: return image, [] # Convert VOC (x1, y1, x2, y2) to COCO (x, y, w, h) person_boxes_coco = person_boxes_voc.copy() person_boxes_coco[:, 2] = person_boxes_voc[:, 2] - person_boxes_voc[:, 0] person_boxes_coco[:, 3] = person_boxes_voc[:, 3] - person_boxes_voc[:, 1] # Step 2: Estimate keypoints with ViTPose pose_inputs = vitpose_processor(image, boxes=[person_boxes_coco], return_tensors="pt").to(device) forward_kwargs = dict(pose_inputs) if VITPOSE_IS_MOE: # COCO dataset index = 0 for vitpose-plus MoE heads forward_kwargs["dataset_index"] = torch.zeros( pose_inputs["pixel_values"].shape[0], dtype=torch.long, device=device ) pose_outputs = vitpose_model(**forward_kwargs) pose_results = vitpose_processor.post_process_pose_estimation(pose_outputs, boxes=[person_boxes_coco]) # Step 3: Build annotations annotations = [] for i, pose_result in enumerate(pose_results[0]): keypoints_xy = pose_result["keypoints"].cpu().numpy() keypoints_scores = pose_result["scores"].cpu().numpy() x1, y1, x2, y2 = person_boxes_voc[i] keypoints = [ { "x": float(keypoints_xy[j][0]), "y": float(keypoints_xy[j][1]), "name": COCO_KEYPOINT_NAMES[j], "confidence": round(float(keypoints_scores[j]), 3), } for j in range(len(keypoints_xy)) ] annotations.append( { "bbox": {"x": float(x1), "y": float(y1), "width": float(x2 - x1), "height": float(y2 - y1)}, "score": round(float(person_scores[i]), 3), "label": "person", "keypoints": keypoints, "connections": SKELETON, } ) return image, annotations, {"score_threshold": (threshold, 1.0)} # ============================================================ # Dispatcher # ============================================================ TASK_OBJECT_DETECTION = "Object Detection" TASK_ZERO_SHOT_DETECTION = "Zero-Shot Detection" TASK_INSTANCE_SEGMENTATION = "Instance Segmentation" TASK_POSE_ESTIMATION = "Pose Estimation" TASK_CHOICES = [ TASK_OBJECT_DETECTION, TASK_ZERO_SHOT_DETECTION, TASK_INSTANCE_SEGMENTATION, TASK_POSE_ESTIMATION, ] _TASK_DISPATCH = { TASK_OBJECT_DETECTION: _detect_rtdetr, TASK_ZERO_SHOT_DETECTION: _detect_gdino, TASK_INSTANCE_SEGMENTATION: _detect_m2f, TASK_POSE_ESTIMATION: _detect_vitpose, } @spaces.GPU @torch.inference_mode() def run_detection( image: Image.Image, task: str, labels: str, threshold: float ) -> tuple[Image.Image, list[dict], dict] | None: if image is None: return None if task == TASK_ZERO_SHOT_DETECTION and not labels.strip(): return None return _TASK_DISPATCH[task](image, labels, threshold) # type: ignore[return-value] # ============================================================ # UI # ============================================================ with gr.Blocks(title="Detection Viewer Demo") as demo: gr.Markdown("# Detection Viewer Demo") gr.Markdown( "Showcase of the **Detection Viewer** Gradio custom component " "— visualizing bounding boxes, segmentation masks, keypoints, and skeleton connections." ) with gr.Row(): with gr.Column(): input_image = gr.Image(label="Input Image", type="pil") task_selector = gr.Radio( choices=TASK_CHOICES, value=TASK_OBJECT_DETECTION, label="Task", ) labels_input = gr.Textbox( label="Labels (comma-separated)", placeholder="person, dog, car, chair", value="person, dog, cat, car", visible=False, ) threshold = gr.Slider(label="Confidence Threshold", minimum=0.0, maximum=1.0, step=0.05, value=0.3) run_btn = gr.Button("Detect", variant="primary") with gr.Column(): viewer = DetectionViewer(label="Detection Results", keypoint_threshold=0.3) task_selector.change( fn=lambda task: gr.update(visible=task == TASK_ZERO_SHOT_DETECTION), inputs=task_selector, outputs=labels_input, ) run_btn.click( fn=run_detection, inputs=[input_image, task_selector, labels_input, threshold], outputs=viewer, ) gr.Examples( examples=[ [str(ASSETS_DIR / "kitchen.jpg"), TASK_OBJECT_DETECTION, "", 0.3], [ str(ASSETS_DIR / "office.jpg"), TASK_ZERO_SHOT_DETECTION, "person, laptop, notebook, pencil, glasses, watch, potted plant, bookshelf, window", 0.3, ], [str(ASSETS_DIR / "traffic.jpg"), TASK_INSTANCE_SEGMENTATION, "", 0.3], [str(ASSETS_DIR / "people.jpg"), TASK_POSE_ESTIMATION, "", 0.3], ], inputs=[input_image, task_selector, labels_input, threshold], fn=run_detection, outputs=viewer, ) if __name__ == "__main__": demo.launch()