Update app.py

app.py

@@ -6,7 +6,6 @@ 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")
 
@@ -18,28 +17,21 @@ MASK_ANNOTATOR = sv.MaskAnnotator()
 LABEL_ANNOTATOR = sv.LabelAnnotator()
 TRACKER = sv.ByteTrack()
 
-def annotate_image(
-    input_image,
-    detections,
-    labels
-) -> np.ndarray:
+def annotate_image(input_image: np.ndarray, detections, labels: List[str]) -> 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
 
-def process_image(
-    input_image,
-    confidence_threshold,
-):
-    results = query(input_image, confidence_threshold)
+def process_image(input_image: np.ndarray, confidence_threshold: float):
+    results = query(Image.fromarray(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
+    return output_image, ", ".join(final_labels)
 
-def query(image, confidence_threshold):
+def query(image: Image.Image, confidence_threshold: float):
     inputs = processor(images=image, return_tensors="pt").to(device)
     with torch.no_grad():
         outputs = model(**inputs)
@@ -48,19 +40,22 @@ def query(image, confidence_threshold):
     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")
+    input_image = gr.Image(label="Input Image", type="numpy")
+    conf = gr.Slider(label="Confidence Threshold", minimum=0.1, maximum=1.0, value=0.6, step=0.05)
+    output_image = gr.Image(label="Output Image", type="numpy")
+    output_text = gr.Textbox(label="Detected Classes")
 
     def process_and_display(input_image, conf):
-        output_img = process_image(input_image, conf)
-        return output_img
+        output_img, detected_classes = process_image(input_image, conf)
+        return output_img, detected_classes
 
     gr.Interface(
         fn=process_and_display,
         inputs=[input_image, conf],
-        outputs=output_image,
+        outputs=[output_image, output_text],
         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,
+        description="This demo uses RT-DETR for object detection in images. Adjust the confidence threshold to see different results.",
     ).launch()
+
+if __name__ == "__main__":
+    run_demo()