Create app.py

app.py

@@ -0,0 +1,67 @@
+from typing import List
+import os
+import numpy as np
+import torch
+import gradio as gr
+from PIL import Image
+from transformers import AutoImageProcessor, AutoModelForObjectDetection
+import supervision as sv
+import spaces
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+processor = AutoImageProcessor.from_pretrained("PekingU/rtdetr_r50vd_coco_o365")
+model = AutoModelForObjectDetection.from_pretrained("PekingU/rtdetr_r50vd_coco_o365").to(device)
+
+BOUNDING_BOX_ANNOTATOR = sv.BoundingBoxAnnotator()
+MASK_ANNOTATOR = sv.MaskAnnotator()
+LABEL_ANNOTATOR = sv.LabelAnnotator()
+TRACKER = sv.ByteTrack()
+
+def annotate_image(
+    input_image,
+    detections,
+    labels
+) -> np.ndarray:
+    output_image = MASK_ANNOTATOR.annotate(input_image, detections)
+    output_image = BOUNDING_BOX_ANNOTATOR.annotate(output_image, detections)
+    output_image = LABEL_ANNOTATOR.annotate(output_image, detections, labels=labels)
+    return output_image
+
+@spaces.GPU
+def process_image(
+    input_image,
+    confidence_threshold,
+):
+    results = query(input_image, confidence_threshold)
+    detections = sv.Detections.from_transformers(results[0])
+    detections = TRACKER.update_with_detections(detections)
+    final_labels = [model.config.id2label[label] for label in detections.class_id.tolist()]
+    output_image = annotate_image(input_image, detections, final_labels)
+    return output_image
+
+def query(image, confidence_threshold):
+    inputs = processor(images=image, return_tensors="pt").to(device)
+    with torch.no_grad():
+        outputs = model(**inputs)
+    target_sizes = torch.tensor([image.size[::-1]])
+    results = processor.post_process_object_detection(outputs=outputs, threshold=confidence_threshold, target_sizes=target_sizes)
+    return results
+
+def run_demo():
+    input_image = gr.inputs.Image(label="Input Image")
+    conf = gr.inputs.Slider(label="Confidence Threshold", minimum=0.1, maximum=1.0, value=0.6, step=0.05)
+    output_image = gr.outputs.Image(label="Output Image")
+
+    def process_and_display(input_image, conf):
+        output_img = process_image(input_image, conf)
+        return output_img
+
+    gr.Interface(
+        fn=process_and_display,
+        inputs=[input_image, conf],
+        outputs=output_image,
+        title="Real Time Object Detection with RT-DETR",
+        description="This demo uses RT-DETR for object detection in images. Adjust the confidence threshold to see different results.",
+        capture_session=True,
+    ).launch()