Update app.py

app.py

@@ -4,51 +4,34 @@ from transformers import pipeline, DetrForObjectDetection, DetrConfig, DetrImage
 import numpy as np
 import cv2
 from PIL import Image
-import warnings
-import logging
-# To suppress all warnings entries
-warnings.filterwarnings('ignore')
-
-# To ignore specific loggings from the Transformers library
-logging.getLogger("transformers").setLevel(logging.ERROR)
-
-def model_is_panoptic(model_name):
-    return "panoptic" in model_name
 
 def load_model(model_name, threshold):
     config = DetrConfig.from_pretrained(model_name, threshold=threshold)
     model = DetrForObjectDetection.from_pretrained(model_name, config=config)
     image_processor = DetrImageProcessor.from_pretrained(model_name)
     return pipeline(task='object-detection', model=model, image_processor=image_processor)
-# Initial model with default threshold
-od_pipe = load_model("facebook/detr-resnet-101", 0.25)
+od_pipe = load_model("facebook/detr-resnet-101", 0.25)  # Default model
 
 def draw_detections(image, detections, model_name):
     np_image = np.array(image)
     np_image = cv2.cvtColor(np_image, cv2.COLOR_RGB2BGR)
     for detection in detections:
-        if model_is_panoptic(model_name):
-            # Handle segmentations for panoptic models
+        if 'mask' in detection:
+            # Interpret and visualize segmentation mask
             mask = detection['mask']
             color = np.random.randint(0, 255, size=3)
             mask = np.round(mask * 255).astype(np.uint8)
             mask = cv2.resize(mask, (image.width, image.height))
             mask_image = np.stack([mask]*3, axis=-1)
             np_image[mask == 255] = np_image[mask == 255] * 0.5 + color * 0.5
-        else:
-            # Handle bounding boxes for standard models
-            score = detection['score']
-            label = detection['label']
+        if 'box' in detection:
+            # Visualize bounding box
             box = detection['box']
-            x_min, y_min = box['xmin'], box['ymin']
-            x_max, y_max = box['xmax'], box['ymax']
+            x_min, y_min, x_max, y_max = [int(b) for b in [box['xmin'], box['ymin'], box['xmax'], box['ymax']]]
             cv2.rectangle(np_image, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
-            label_text = f'{label} {score:.2f}'
-            cv2.putText(np_image, label_text, (x_min, y_min - 10), cv2.FONT_HERSHEY_SIMPLEX, 1.5, (255, 255, 255), 4)
     final_image = cv2.cvtColor(np_image, cv2.COLOR_BGR2RGB)
-    final_pil_image = Image.fromarray(final_image)
-    return final_pil_image
-    
+    return Image.fromarray(final_image)
+
 def get_pipeline_prediction(model_name, threshold, pil_image):
     global od_pipe
     od_pipe = load_model(model_name, threshold)
@@ -56,6 +39,7 @@ def get_pipeline_prediction(model_name, threshold, pil_image):
         if not isinstance(pil_image, Image.Image):
             pil_image = Image.fromarray(np.array(pil_image).astype('uint8'), 'RGB')
         result = od_pipe(pil_image)
+        print("Detection Output:", result)  # Debug: Check the output structure
         processed_image = draw_detections(pil_image, result, model_name)
         description = f'Model used: {model_name}, Detection Threshold: {threshold}'
         return processed_image, result, description
@@ -77,5 +61,4 @@ with gr.Blocks() as demo:
             with gr.Tab("Description"):
                 description_output = gr.Textbox()
     run_button.click(get_pipeline_prediction, inputs=[model_dropdown, threshold_slider, inp_image], outputs=[output_image, output_data, description_output])
-    
 demo.launch()