Create app.py

app.py

@@ -0,0 +1,568 @@
+import os
+import cv2
+import tqdm
+import uuid
+import logging
+
+import torch
+import spaces
+import trackers
+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)
+
+class TrackingAlgorithm:
+    BYTETRACK = "ByteTrack (2021)"
+    DEEPSORT = "DeepSORT (2017)"
+    SORT = "SORT (2016)"
+
+TRACKERS = [None, TrackingAlgorithm.BYTETRACK, TrackingAlgorithm.DEEPSORT, TrackingAlgorithm.SORT]
+VIDEO_EXAMPLES = [
+    {"path": "./examples/videos/dogs_running.mp4", "label": "Local Video", "tracker": None, "classes": "all"},
+    {"path": "./examples/videos/traffic.mp4", "label": "Local Video", "tracker": TrackingAlgorithm.BYTETRACK, "classes": "car, truck, bus"},
+    {"path": "./examples/videos/fast_and_furious.mp4", "label": "Local Video", "tracker": None, "classes": "all"},
+    {"path": "./examples/videos/break_dance.mp4", "label": "Local Video", "tracker": None, "classes": "all"},
+]
+
+
+# Create a color palette for visualization
+# These hex color codes define different colors for tracking different objects
+color = sv.ColorPalette.from_hex([
+    "#ffff00", "#ff9b00", "#ff8080", "#ff66b2", "#ff66ff", "#b266ff",
+    "#9999ff", "#3399ff", "#66ffff", "#33ff99", "#66ff66", "#99ff00"
+])
+
+
+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,
+    classes: Optional[List[str]] = None,
+):
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model, image_processor = get_model_and_processor(checkpoint)
+    model = model.to(device)
+
+    if classes is not None:
+        wrong_classes = [cls for cls in classes if cls not in model.config.label2id]
+        if wrong_classes:
+            gr.Warning(f"Classes not found in model config: {wrong_classes}")
+        keep_ids = [model.config.label2id[cls] for cls in classes if cls in model.config.label2id]
+    else:
+        keep_ids = None
+
+    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()}
+        if keep_ids is not None:
+            keep = torch.isin(results[i]["labels"], torch.tensor(keep_ids))
+            results[i] = {k: v[keep] for k, v in results[i].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_height, image_width 
+    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 get_tracker(tracker: str, fps: float):
+    if tracker == TrackingAlgorithm.SORT:
+        return trackers.SORTTracker(frame_rate=fps)
+    elif tracker == TrackingAlgorithm.DEEPSORT:
+        feature_extractor = trackers.DeepSORTFeatureExtractor.from_timm("mobilenetv4_conv_small.e1200_r224_in1k", device="cpu")
+        return trackers.DeepSORTTracker(feature_extractor, frame_rate=fps)
+    elif tracker == TrackingAlgorithm.BYTETRACK:  
+        return sv.ByteTrack(frame_rate=int(fps))
+    else:
+        raise ValueError(f"Invalid tracker: {tracker}")
+
+
+def update_tracker(tracker, detections, frame):
+    tracker_name = tracker.__class__.__name__
+    if tracker_name == "SORTTracker":
+        return tracker.update(detections)
+    elif tracker_name == "DeepSORTTracker":
+        return tracker.update(detections, frame)
+    elif tracker_name == "ByteTrack":
+        return tracker.update_with_detections(detections)
+    else:
+        raise ValueError(f"Invalid tracker: {tracker}")
+
+
+def process_video(
+    video_path: str,
+    checkpoint: str,
+    tracker_algorithm: Optional[str] = None,
+    classes: str = "all",
+    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 = max(1, 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]
+
+    # Set the color lookup mode to assign colors by track ID
+    # This mean objects with the same track ID will be annotated by the same color
+    color_lookup = sv.ColorLookup.TRACK if tracker_algorithm else sv.ColorLookup.CLASS
+
+    box_annotator = sv.BoxAnnotator(color, color_lookup=color_lookup, thickness=1)
+    label_annotator = sv.LabelAnnotator(color, color_lookup=color_lookup, text_scale=0.5)
+    trace_annotator = sv.TraceAnnotator(color, color_lookup=color_lookup, thickness=1, trace_length=100)
+
+    # preprocess classes
+    if classes != "all":
+        classes_list = [cls.strip().lower() for cls in classes.split(",")]
+    else:
+        classes_list = None
+
+    results, id2label = detect_objects(
+        images=np.array(frames),
+        checkpoint=checkpoint,
+        confidence_threshold=confidence_threshold,
+        target_size=(target_height, target_width),
+        classes=classes_list,
+    )
+
+
+    annotated_frames = []
+
+    # detections
+    if tracker_algorithm:
+        tracker = get_tracker(tracker_algorithm, target_fps)
+        for frame, result in progress.tqdm(zip(frames, results), desc="Tracking objects", total=len(frames)):
+            detections = sv.Detections.from_transformers(result, id2label=id2label)
+            detections = detections.with_nms(threshold=0.95, class_agnostic=True)
+            detections = update_tracker(tracker, detections, frame)
+            labels = [f"#{tracker_id} {id2label[class_id]}" for class_id, tracker_id in zip(detections.class_id, detections.tracker_id)]
+            annotated_frame = box_annotator.annotate(scene=frame, detections=detections)
+            annotated_frame = label_annotator.annotate(scene=annotated_frame, detections=detections, labels=labels)
+            annotated_frame = trace_annotator.annotate(scene=annotated_frame, detections=detections)
+            annotated_frames.append(annotated_frame)
+        
+    else:
+        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.Dropdown(
+            choices=TRACKERS,
+            label="Select Tracker (Optional)",
+            value=None,
+            elem_classes="input-component",
+        ),
+        gr.TextArea(
+            label="Specify Class Names to Detect (comma separated)",
+            value="all",
+            lines=1,
+            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](https://huggingface.co/docs/transformers/main/model_doc/d_fine) models.
+         - **Image** and **Video** modes are supported.
+         - Select a model and adjust the confidence threshold to see detections!
+         - On video mode, you can enable tracking powered by [Supervision](https://github.com/roboflow/supervision) and [Trackers](https://github.com/roboflow/trackers) from Roboflow.
+        """,
+        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_tracker, video_classes, 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, example["tracker"], example["classes"], DEFAULT_CONFIDENCE_THRESHOLD]
+                    for example in VIDEO_EXAMPLES
+                ],
+                inputs=[video_input, video_checkpoint, video_tracker, video_classes, 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,
+            None,
+            "all",
+            DEFAULT_CONFIDENCE_THRESHOLD,
+            None,
+        ),
+        outputs=[
+            video_input,
+            video_checkpoint,
+            video_tracker,
+            video_classes,
+            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_tracker, video_classes, video_confidence_threshold],
+        outputs=[video_output],
+    )
+
+if __name__ == "__main__":
+    demo.queue(max_size=20).launch()