Update app.py

app.py

@@ -1,34 +1,57 @@
+python
 import os
 import gradio as gr
 from transformers import pipeline, DetrForObjectDetection, DetrConfig, DetrImageProcessor
+import numpy as np
+import cv2
+from PIL import Image
 
-# Initialize the configuration for DetrForObjectDetection
-config = DetrConfig.from_pretrained("facebook/detr-resnet-50")
+def draw_detections(image, detections):
+    # Convert PIL image to a numpy array
+    np_image = np.array(image)
+    
+    # Convert RGB to BGR for OpenCV
+    np_image = cv2.cvtColor(np_image, cv2.COLOR_RGB2BGR)
+    
+    for detection in detections:
+        # Each detection includes ['score', 'label', 'box']
+        score = detection['score']
+        label = detection['label']
+        box = detection['box']
+        x_min, y_min, x_max, y_max = map(int, box)
+        cv2.rectangle(np_image, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
+        cv2.putText(np_image, f'{label} {score:.2f}', (x_min, max(y_min - 10, 0)),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
+    
+    # Convert BGR to RGB for displaying
+    final_image = cv2.cvtColor(np_image, cv2.COLOR_BGR2RGB)
+    # Convert the numpy array to PIL Image
+    final_pil_image = Image.fromarray(final_image)
+    return final_pil_image
 
-# Create the model for object detection using the specified configuration
+# Initialize objects from transformers
+config = DetrConfig.from_pretrained("facebook/detr-resnet-50")
 model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50", config=config)
-
-# Initialize the image processor for DETR
 image_processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
-
-# Initialize the object detection pipeline with the model and image processor
 od_pipe = pipeline(task='object-detection', model=model, image_processor=image_processor)
 
 def get_pipeline_prediction(pil_image):
-    # Run the object detection pipeline on the input image
+    # Run the object detection pipeline
     pipeline_output = od_pipe(pil_image)
     
-    # You might need to implement or adjust the rendering function based on the `pipeline_output`
-    # The `render_results_in_image` function is assumed here to draw bounding boxes and labels on the input image,
-    # but you'll need to define it according to your specific needs.
-    # For now, the output is directly returned since the question doesn't define `render_results_in_image`.
-    return pipeline_output
+    # Draw the detection results on the image
+    processed_image = draw_detections(pil_image, pipeline_output)
+    
+    # Provide both the image and the JSON detection results
+    return processed_image, pipeline_output
 
 demo = gr.Interface(
-  fn=get_pipeline_prediction,
-  inputs=gr.Image(label="Input image", 
-                  type="pil"),
-  outputs=gr.JSON(label="Detected objects")  # Adjusted to show JSON output if rendering function is not defined
+    fn=get_pipeline_prediction,
+    inputs=gr.Image(label="Input image", type="pil"),
+    outputs=[
+        gr.Image(label="Annotated Image"),
+        gr.JSON(label="Detected Objects")
+    ]
 )
 
-demo.launch()
+demo.launch()