Update app.py

app.py

@@ -11,41 +11,62 @@ model = YOLO('best.pt')
 # Define class names based on your model
 class_names = ['null', 'pollen', 'queen', 'queen_cell', 'varroa']
 
-def predict_bee(image):
+def predict_bee(image, confidence_threshold=0.25):
     """
     Predict bee type and health status from uploaded image
     """
     try:
+        # Debug: Print image info
+        print(f"Input image size: {image.size}")
+        print(f"Input image mode: {image.mode}")
+        print(f"Confidence threshold: {confidence_threshold}")
+        
         # Convert PIL image to numpy array
         img_array = np.array(image)
+        print(f"Array shape: {img_array.shape}")
+        
+        # Ensure image is in RGB format
+        if len(img_array.shape) == 3 and img_array.shape[2] == 3:
+            # Convert RGB to BGR for OpenCV compatibility
+            img_array = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
         
-        # Run inference
-        results = model(img_array)
+        # Run inference with explicit parameters
+        results = model(img_array, conf=confidence_threshold, iou=0.45, verbose=True)
         
         # Process results
         predictions = []
-        annotated_image = img_array.copy()
+        annotated_image = cv2.cvtColor(img_array, cv2.COLOR_BGR2RGB)  # Convert back to RGB for display
         
-        for result in results:
+        print(f"Number of results: {len(results)}")
+        
+        for i, result in enumerate(results):
+            print(f"Processing result {i}")
             # Get boxes, scores, and classes
             boxes = result.boxes
             
+            print(f"Boxes object: {boxes}")
             if boxes is not None:
-                for box in boxes:
+                print(f"Number of detections: {len(boxes)}")
+                
+                for j, box in enumerate(boxes):
+                    print(f"Processing box {j}")
+                    
                     # Get coordinates
                     x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)
                     
                     # Get confidence and class
-                    confidence = box.conf[0].cpu().numpy()
+                    confidence = float(box.conf[0].cpu().numpy())
                     class_id = int(box.cls[0].cpu().numpy())
                     
+                    print(f"Detection: class_id={class_id}, confidence={confidence:.3f}, bbox=[{x1},{y1},{x2},{y2}]")
+                    
                     # Get class name
                     class_name = class_names[class_id] if class_id < len(class_names) else f"Class_{class_id}"
                     
                     # Add to predictions
                     predictions.append({
                         'class': class_name,
-                        'confidence': float(confidence),
+                        'confidence': confidence,
                         'bbox': [int(x1), int(y1), int(x2), int(y2)]
                     })
                     
@@ -60,6 +81,10 @@ def predict_bee(image):
                                 (x1 + label_size[0], y1), color, -1)
                     cv2.putText(annotated_image, label, (x1, y1 - 5), 
                               cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
+            else:
+                print("No boxes detected")
+        
+        print(f"Total predictions: {len(predictions)}")
         
         # Convert back to PIL Image
         annotated_image = Image.fromarray(annotated_image)
@@ -127,6 +152,15 @@ def create_interface():
                     height=400
                 )
                 
+                confidence_slider = gr.Slider(
+                    minimum=0.1,
+                    maximum=0.9,
+                    value=0.25,
+                    step=0.05,
+                    label="Confidence Threshold",
+                    info="Lower values = more detections (but may include false positives)"
+                )
+                
                 submit_btn = gr.Button(
                     "🔍 Analyze Hive", 
                     variant="primary",
@@ -173,13 +207,19 @@ def create_interface():
         # Event handlers
         submit_btn.click(
             fn=predict_bee,
-            inputs=image_input,
+            inputs=[image_input, confidence_slider],
             outputs=[image_output, text_output]
         )
         
         image_input.change(
             fn=predict_bee,
-            inputs=image_input,
+            inputs=[image_input, confidence_slider],
+            outputs=[image_output, text_output]
+        )
+        
+        confidence_slider.change(
+            fn=predict_bee,
+            inputs=[image_input, confidence_slider],
             outputs=[image_output, text_output]
         )