Update app.py

app.py

@@ -11,9 +11,8 @@ model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50", config
 image_processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
 
 def load_model(threshold):
-    # Since changing threshold at runtime for models isn't typically supported directly by the transformers pipeline,
-    # we reinitialize the model with the desired configuration when needed.
-    config = DetrConfig.from_pretrained("facebook/detr-resnet-50", num_labels=91, threshold=threshold)
+    # Reinitialize the model with the desired detection threshold
+    config = DetrConfig.from_pretrained("facebook/detr-resnet-50", threshold=threshold)
     model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50", config=config)
     image_processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
     return pipeline(task='object-detection', model=model, image_processor=image_processor)
@@ -27,6 +26,7 @@ def draw_detections(image, detections):
     # Convert RGB to BGR for OpenCV
     np_image = cv2.cvtColor(np_image, cv2.COLOR_RGB2BGR)
     
+    # Draw detections
     for detection in detections:
         score = detection['score']
         label = detection['label']
@@ -36,10 +36,10 @@ def draw_detections(image, detections):
         x_max = box['xmax']
         y_max = box['ymax']
 
-        # Draw rectangles and text with a larger font
+        # Increase font size for better visibility
         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, 0.8, (255, 255, 255), 2)
+        cv2.putText(np_image, label_text, (x_min, y_min - 10), cv2.FONT_HERSHEY_SIMPLEX, 1.5, (255, 255, 255), 4)
     
     # Convert BGR to RGB for displaying
     final_image = cv2.cvtColor(np_image, cv2.COLOR_BGR2RGB)
@@ -48,30 +48,39 @@ def draw_detections(image, detections):
 
 def get_pipeline_prediction(threshold, pil_image):
     global od_pipe
-    if od_pipe.config.threshold != threshold:
-        od_pipe = load_model(threshold)
     try:
-        pil_image = Image.fromarray(np.array(pil_image))
+        # Check if the model threshold needs adjusting
+        if od_pipe.config.threshold != threshold:
+            od_pipe = load_model(threshold)
+            print("Model reloaded with new threshold:", threshold)
+        
+        # Ensure input is a PIL image
+        if not isinstance(pil_image, Image.Image):
+            pil_image = Image.fromarray(np.array(pil_image).astype('uint8'), 'RGB')
+        
+        # Run detection and return annotated image and results
         pipeline_output = od_pipe(pil_image)
         processed_image = draw_detections(pil_image, pipeline_output)
         return processed_image, pipeline_output
     except Exception as e:
-        print(f"An error occurred: {str(e)}")
-        return pil_image, {"error": str(e)}
+        error_message = f"An error occurred: {str(e)}"
+        print(error_message)
+        return pil_image, {"error": error_message}
 
-# Define the Gradio blocks interface
+# Gradio interface
 with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             inp_image = gr.Image(label="Input image")
-            slider = gr.Slider(minimum=0, maximum=1, step=0.05, label="Adjust Detection Sensitivity", value=0.5)
-            gr.Markdown("Adjust the slider to change the detection sensitivity.")
+            slider = gr.Slider(minimum=0, maximum=1, step=0.05, label="Detection Sensitivity", value=0.5)
+            gr.Markdown("Adjust the slider to change detection sensitivity.")
             btn_run = gr.Button('Run Detection')
         with gr.Column():
             with gr.Tab("Annotated Image"):
                 out_image = gr.Image()
             with gr.Tab("Detection Results"):
                 out_json = gr.JSON()
+
     btn_run.click(get_pipeline_prediction, inputs=[slider, inp_image], outputs=[out_image, out_json])
 
 demo.launch()