import os import cv2 import tqdm import uuid import logging import torch import spaces import numpy as np import gradio as gr import imageio.v3 as iio import supervision as sv from pathlib import Path from functools import lru_cache from typing import List, Optional, Tuple from PIL import Image from transformers import AutoModelForObjectDetection, AutoImageProcessor from transformers.image_utils import load_image # Configuration constants CHECKPOINTS = [ "ustc-community/dfine-medium-obj2coco", "ustc-community/dfine-medium-coco", "ustc-community/dfine-medium-obj365", "ustc-community/dfine-nano-coco", "ustc-community/dfine-small-coco", "ustc-community/dfine-large-coco", "ustc-community/dfine-xlarge-coco", "ustc-community/dfine-small-obj365", "ustc-community/dfine-large-obj365", "ustc-community/dfine-xlarge-obj365", "ustc-community/dfine-small-obj2coco", "ustc-community/dfine-large-obj2coco-e25", "ustc-community/dfine-xlarge-obj2coco", ] DEFAULT_CHECKPOINT = CHECKPOINTS[0] DEFAULT_CONFIDENCE_THRESHOLD = 0.3 TORCH_DTYPE = torch.float32 # Image IMAGE_EXAMPLES = [ {"path": "./examples/images/tennis.jpg", "use_url": False, "url": "", "label": "Local Image"}, {"path": "./examples/images/dogs.jpg", "use_url": False, "url": "", "label": "Local Image"}, {"path": "./examples/images/nascar.jpg", "use_url": False, "url": "", "label": "Local Image"}, {"path": "./examples/images/crossroad.jpg", "use_url": False, "url": "", "label": "Local Image"}, { "path": None, "use_url": True, "url": "https://live.staticflickr.com/65535/33021460783_1646d43c54_b.jpg", "label": "Flickr Image", }, ] # Video MAX_NUM_FRAMES = 250 BATCH_SIZE = 4 ALLOWED_VIDEO_EXTENSIONS = {".mp4", ".avi", ".mov"} VIDEO_OUTPUT_DIR = Path("static/videos") VIDEO_OUTPUT_DIR.mkdir(parents=True, exist_ok=True) VIDEO_EXAMPLES = [ {"path": "./examples/videos/dogs_running.mp4", "label": "Local Video"}, {"path": "./examples/videos/traffic.mp4", "label": "Local Video"}, {"path": "./examples/videos/fast_and_furious.mp4", "label": "Local Video"}, {"path": "./examples/videos/break_dance.mp4", "label": "Local Video"}, ] logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" ) logger = logging.getLogger(__name__) @lru_cache(maxsize=3) def get_model_and_processor(checkpoint: str): model = AutoModelForObjectDetection.from_pretrained(checkpoint, torch_dtype=TORCH_DTYPE) image_processor = AutoImageProcessor.from_pretrained(checkpoint) return model, image_processor @spaces.GPU(duration=20) def detect_objects( checkpoint: str, images: List[np.ndarray] | np.ndarray, confidence_threshold: float = DEFAULT_CONFIDENCE_THRESHOLD, target_size: Optional[Tuple[int, int]] = None, batch_size: int = BATCH_SIZE, ): device = "cuda" if torch.cuda.is_available() else "cpu" model, image_processor = get_model_and_processor(checkpoint) model = model.to(device) if isinstance(images, np.ndarray) and images.ndim == 4: images = [x for x in images] # split video array into list of images batches = [images[i:i + batch_size] for i in range(0, len(images), batch_size)] results = [] for batch in tqdm.tqdm(batches, desc="Processing frames"): # preprocess images inputs = image_processor(images=batch, return_tensors="pt") inputs = inputs.to(device).to(TORCH_DTYPE) # forward pass with torch.no_grad(): outputs = model(**inputs) # postprocess outputs if target_size: target_sizes = [target_size] * len(batch) else: target_sizes = [(image.shape[0], image.shape[1]) for image in batch] batch_results = image_processor.post_process_object_detection( outputs, target_sizes=target_sizes, threshold=confidence_threshold ) results.extend(batch_results) # move results to cpu for i, result in enumerate(results): results[i] = {k: v.cpu() for k, v in result.items()} return results, model.config.id2label def process_image( checkpoint: str = DEFAULT_CHECKPOINT, image: Optional[Image.Image] = None, url: Optional[str] = None, use_url: bool = False, confidence_threshold: float = DEFAULT_CONFIDENCE_THRESHOLD, ): if not use_url: url = None if (image is None) ^ bool(url): raise ValueError(f"Either image or url must be provided, but not both.") if url: image = load_image(url) results, id2label = detect_objects( checkpoint=checkpoint, images=[np.array(image)], confidence_threshold=confidence_threshold, ) result = results[0] # first image in batch (we have batch size 1) annotations = [] for label, score, box in zip(result["labels"], result["scores"], result["boxes"]): text_label = id2label[label.item()] formatted_label = f"{text_label} ({score:.2f})" x_min, y_min, x_max, y_max = box.cpu().numpy().round().astype(int) x_min = max(0, x_min) y_min = max(0, y_min) x_max = min(image.width - 1, x_max) y_max = min(image.height - 1, y_max) annotations.append(((x_min, y_min, x_max, y_max), formatted_label)) return (image, annotations) def get_target_size(image_height, image_width, max_size: int): if image_height < max_size and image_width < max_size: new_height, new_width = image_width, image_height elif image_height > image_width: new_height = max_size new_width = int(image_width * max_size / image_height) else: new_width = max_size new_height = int(image_height * max_size / image_width) # make even (for video codec compatibility) new_height = new_height // 2 * 2 new_width = new_width // 2 * 2 return new_width, new_height def read_video_k_frames(video_path: str, k: int, read_every_i_frame: int = 1): cap = cv2.VideoCapture(video_path) frames = [] i = 0 progress_bar = tqdm.tqdm(total=k, desc="Reading frames") while cap.isOpened() and len(frames) < k: ret, frame = cap.read() if not ret: break if i % read_every_i_frame == 0: frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) progress_bar.update(1) i += 1 cap.release() progress_bar.close() return frames def process_video( video_path: str, checkpoint: str, confidence_threshold: float = DEFAULT_CONFIDENCE_THRESHOLD, progress: gr.Progress = gr.Progress(track_tqdm=True), ) -> str: if not video_path or not os.path.isfile(video_path): raise ValueError(f"Invalid video path: {video_path}") ext = os.path.splitext(video_path)[1].lower() if ext not in ALLOWED_VIDEO_EXTENSIONS: raise ValueError(f"Unsupported video format: {ext}, supported formats: {ALLOWED_VIDEO_EXTENSIONS}") video_info = sv.VideoInfo.from_video_path(video_path) read_each_i_frame = video_info.fps // 25 target_fps = video_info.fps / read_each_i_frame target_width, target_height = get_target_size(video_info.height, video_info.width, 1080) n_frames_to_read = min(MAX_NUM_FRAMES, video_info.total_frames // read_each_i_frame) frames = read_video_k_frames(video_path, n_frames_to_read, read_each_i_frame) frames = [cv2.resize(frame, (target_width, target_height), interpolation=cv2.INTER_CUBIC) for frame in frames] box_annotator = sv.BoxAnnotator(thickness=1) label_annotator = sv.LabelAnnotator(text_scale=0.5) results, id2label = detect_objects( images=np.array(frames), checkpoint=checkpoint, confidence_threshold=confidence_threshold, target_size=(target_height, target_width), ) annotated_frames = [] for frame, result in tqdm.tqdm(zip(frames, results), desc="Annotating frames", total=len(frames)): detections = sv.Detections.from_transformers(result, id2label=id2label) detections = detections.with_nms(threshold=0.95, class_agnostic=True) annotated_frame = box_annotator.annotate(scene=frame, detections=detections) annotated_frame = label_annotator.annotate(scene=annotated_frame, detections=detections) annotated_frames.append(annotated_frame) output_filename = os.path.join(VIDEO_OUTPUT_DIR, f"output_{uuid.uuid4()}.mp4") iio.imwrite(output_filename, annotated_frames, fps=target_fps, codec="h264") return output_filename def create_image_inputs() -> List[gr.components.Component]: return [ gr.Image( label="Upload Image", type="pil", sources=["upload", "webcam"], interactive=True, elem_classes="input-component", ), gr.Checkbox(label="Use Image URL Instead", value=False), gr.Textbox( label="Image URL", placeholder="https://example.com/image.jpg", visible=False, elem_classes="input-component", ), gr.Dropdown( choices=CHECKPOINTS, label="Select Model Checkpoint", value=DEFAULT_CHECKPOINT, elem_classes="input-component", ), gr.Slider( minimum=0.1, maximum=1.0, value=DEFAULT_CONFIDENCE_THRESHOLD, step=0.1, label="Confidence Threshold", elem_classes="input-component", ), ] def create_video_inputs() -> List[gr.components.Component]: return [ gr.Video( label="Upload Video", sources=["upload"], interactive=True, format="mp4", # Ensure MP4 format elem_classes="input-component", ), gr.Dropdown( choices=CHECKPOINTS, label="Select Model Checkpoint", value=DEFAULT_CHECKPOINT, elem_classes="input-component", ), gr.Slider( minimum=0.1, maximum=1.0, value=DEFAULT_CONFIDENCE_THRESHOLD, step=0.1, label="Confidence Threshold", elem_classes="input-component", ), ] def create_button_row() -> List[gr.Button]: return [ gr.Button( f"Detect Objects", variant="primary", elem_classes="action-button" ), gr.Button(f"Clear", variant="secondary", elem_classes="action-button"), ] # Gradio interface with gr.Blocks(theme=gr.themes.Ocean()) as demo: gr.Markdown( """ # Object Detection Demo Experience state-of-the-art object detection with USTC's D-Fine models. - **Image** and **Video** modes are supported. - Select a model and adjust the confidence threshold to see detections! """, elem_classes="header-text", ) with gr.Tabs(): with gr.Tab("Image"): with gr.Row(): with gr.Column(scale=1, min_width=300): with gr.Group(): ( image_input, use_url, url_input, image_model_checkpoint, image_confidence_threshold, ) = create_image_inputs() image_detect_button, image_clear_button = create_button_row() with gr.Column(scale=2): image_output = gr.AnnotatedImage( label="Detection Results", show_label=True, color_map=None, elem_classes="output-component", ) gr.Examples( examples=[ [ DEFAULT_CHECKPOINT, example["path"], example["url"], example["use_url"], DEFAULT_CONFIDENCE_THRESHOLD, ] for example in IMAGE_EXAMPLES ], inputs=[ image_model_checkpoint, image_input, url_input, use_url, image_confidence_threshold, ], outputs=[image_output], fn=process_image, label="Select an image example to populate inputs", cache_examples=True, cache_mode="lazy", ) with gr.Tab("Video"): gr.Markdown( f"The input video will be processed in ~25 FPS (up to {MAX_NUM_FRAMES} frames in result)." ) with gr.Row(): with gr.Column(scale=1, min_width=300): with gr.Group(): video_input, video_checkpoint, video_confidence_threshold = create_video_inputs() video_detect_button, video_clear_button = create_button_row() with gr.Column(scale=2): video_output = gr.Video( label="Detection Results", format="mp4", # Explicit MP4 format elem_classes="output-component", ) gr.Examples( examples=[ [example["path"], DEFAULT_CHECKPOINT, DEFAULT_CONFIDENCE_THRESHOLD] for example in VIDEO_EXAMPLES ], inputs=[video_input, video_checkpoint, video_confidence_threshold], outputs=[video_output], fn=process_video, cache_examples=False, label="Select a video example to populate inputs", ) # Dynamic visibility for URL input use_url.change( fn=lambda x: gr.update(visible=x), inputs=use_url, outputs=url_input, ) # Image clear button image_clear_button.click( fn=lambda: ( None, False, "", DEFAULT_CHECKPOINT, DEFAULT_CONFIDENCE_THRESHOLD, None, ), outputs=[ image_input, use_url, url_input, image_model_checkpoint, image_confidence_threshold, image_output, ], ) # Video clear button video_clear_button.click( fn=lambda: ( None, DEFAULT_CHECKPOINT, DEFAULT_CONFIDENCE_THRESHOLD, None, ), outputs=[ video_input, video_checkpoint, video_confidence_threshold, video_output, ], ) # Image detect button image_detect_button.click( fn=process_image, inputs=[ image_model_checkpoint, image_input, url_input, use_url, image_confidence_threshold, ], outputs=[image_output], ) # Video detect button video_detect_button.click( fn=process_video, inputs=[video_input, video_checkpoint, video_confidence_threshold], outputs=[video_output], ) if __name__ == "__main__": demo.queue(max_size=20).launch()