Update app.py

app.py

@@ -16,7 +16,6 @@ CLASS_NAMES = [
     "card_oracle_text", "card_power_toughness"
 ]
 
-# Define inference function
 def segment_card(image):
     image = np.array(image)  # Convert PIL image to NumPy array
     results = model(image)   # Run YOLO inference
@@ -24,45 +23,57 @@ def segment_card(image):
     # Convert to OpenCV format
     annotated_image = image.copy()
 
+    # Dictionary to store the highest confidence detection for each class
+    best_detections = {}
+
     # Extract bounding boxes and labels
     for result in results:
         for box in result.boxes:
             x1, y1, x2, y2 = map(int, box.xyxy[0])  # Bounding box coordinates
             class_id = int(box.cls[0])  # Class index
-            label = CLASS_NAMES[class_id]  # Get class label
-            
-            # Draw bounding box
-            cv2.rectangle(annotated_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
-            
-            # Draw label text with background
-            font = cv2.FONT_HERSHEY_SIMPLEX
-            font_scale = 0.5
-            font_thickness = 2
-            text_size = cv2.getTextSize(label, font, font_scale, font_thickness)[0]
-            text_x, text_y = x1, y1 - 10
+            conf = float(box.conf[0])  # Confidence score
             
-            # Ensure text doesn't go out of bounds
-            text_y = max(text_y, text_size[1] + 10)
-            
-            # Draw filled rectangle for text background
-            cv2.rectangle(
-                annotated_image, 
-                (text_x, text_y - text_size[1] - 5), 
-                (text_x + text_size[0] + 5, text_y + 5), 
-                (0, 255, 0), 
-                -1
-            )
-            
-            # Put text label on the image
-            cv2.putText(
-                annotated_image, 
-                label, 
-                (text_x, text_y), 
-                font, 
-                font_scale, 
-                (0, 0, 0),  # Text color (black for contrast)
-                font_thickness
-            )
+            # Check if we have seen this class before
+            if class_id not in best_detections or conf > best_detections[class_id]["conf"]:
+                best_detections[class_id] = {"bbox": (x1, y1, x2, y2), "conf": conf}
+
+    # Draw only the best bounding box for each class
+    for class_id, data in best_detections.items():
+        x1, y1, x2, y2 = data["bbox"]
+        label = CLASS_NAMES[class_id]
+
+        # Draw bounding box
+        cv2.rectangle(annotated_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
+
+        # Draw label text with background
+        font = cv2.FONT_HERSHEY_SIMPLEX
+        font_scale = 0.5
+        font_thickness = 2
+        text_size = cv2.getTextSize(label, font, font_scale, font_thickness)[0]
+        text_x, text_y = x1, y1 - 10
+
+        # Ensure text doesn't go out of bounds
+        text_y = max(text_y, text_size[1] + 10)
+
+        # Draw filled rectangle for text background
+        cv2.rectangle(
+            annotated_image, 
+            (text_x, text_y - text_size[1] - 5), 
+            (text_x + text_size[0] + 5, text_y + 5), 
+            (0, 255, 0), 
+            -1
+        )
+
+        # Put text label on the image
+        cv2.putText(
+            annotated_image, 
+            label, 
+            (text_x, text_y), 
+            font, 
+            font_scale, 
+            (0, 0, 0),  # Text color (black for contrast)
+            font_thickness
+        )
 
     return Image.fromarray(annotated_image)  # Convert back to PIL Image
 
@@ -76,4 +87,4 @@ iface = gr.Interface(
 )
 
 # Launch the app
-iface.launch()
+iface.launch()