Update app.py

app.py

@@ -7,20 +7,28 @@ from ultralytics import YOLO
 duck_model = YOLO('https://huggingface.co/brainwavecollective/yolo8n-rubber-duck-detector/resolve/main/yolov8n_rubberducks.pt')
 standard_model = YOLO('yolov8n.pt')
 
-def process_image(image, model):
+def process_image(image, model, is_standard_model=True):
     results = model(image)
     processed_image = image.copy()
     
     for r in results:
         boxes = r.boxes
         for box in boxes:
-            x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy())
-            conf = float(box.conf[0])
             cls = int(box.cls[0])
             class_name = model.names[cls]
             
+            # For standard model, only show teddy bears
+            if is_standard_model and class_name != "teddy bear":
+                continue
+                
+            x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy())
+            conf = float(box.conf[0])
+            
+            # Rename class to "rubber duck"
+            display_name = "rubber duck" if not is_standard_model else class_name
+            
             cv2.rectangle(processed_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
-            label = f"{class_name} ({conf:.2f})"
+            label = f"{display_name} ({conf:.2f})"
             cv2.putText(processed_image, label, (x1, y1-10), 
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
     
@@ -29,8 +37,8 @@ def process_image(image, model):
 def compare_models(input_image):
     image = np.array(input_image)
     
-    standard_image = process_image(image, standard_model)
-    duck_image = process_image(image, duck_model)
+    standard_image = process_image(image, standard_model, is_standard_model=True)
+    duck_image = process_image(image, duck_model, is_standard_model=False)
     
     height, width = image.shape[:2]
     gap = 20  # Add a 20-pixel gap between images