Update app.py

app.py

@@ -30,19 +30,52 @@ def detect_objects(image):
 
     return image
 
+def detect_labels(image):
+    # Load the pre-trained DETR model
+    processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
+    model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
+
+    inputs = processor(images=image, return_tensors="pt")
+    outputs = model(**inputs)
 
-def upload_image(file):
+    # convert outputs (bounding boxes and class logits) to COCO API
+    # let's only keep detections with score > 0.9
+    target_sizes = torch.tensor([image.size[::-1]])
+    results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]
+
+    labels = []
+    for label_id in results["labels"]:
+        labels.append(model.config.id2label[label_id.item()])
+
+    return labels
+
+def upload_image_with_boxes(file):
     image = Image.open(file.name)
     image_with_boxes = detect_objects(image)
     return image_with_boxes
 
-iface = gr.Interface(
-    fn=upload_image,
+def upload_image_with_labels(file):
+    image = Image.open(file.name)
+    labels = detect_labels(image)
+    return ", ".join(labels)
+
+iface_boxes = gr.Interface(
+    fn=upload_image_with_boxes,
     inputs="file",
     outputs="image",
-    title="Object Detection",
-    description="Upload an image and detect objects using DETR model.",
+    title="Object Detection with Boxes",
+    description="Upload an image and detect objects using DETR model. Boxes will be drawn around the detected objects.",
+    allow_flagging=False
+)
+
+iface_labels = gr.Interface(
+    fn=upload_image_with_labels,
+    inputs="file",
+    outputs="text",
+    title="Detected Object Labels",
+    description="Upload an image and get the detected object labels using DETR model.",
     allow_flagging=False
 )
 
-iface.launch()
+iface_boxes.launch()
+iface_labels.launch()