Update app.py

app.py

@@ -9,13 +9,14 @@ from PIL import Image
 model_path = "best.pt"
 model = YOLO(model_path)
 
-# Class names (Ensure these match your YOLO class labels)
+# Class names
 CLASS_NAMES = [
     "card_title", "card_art", "card_type", 
     "card_set_symbol", "card_mana_cost", 
     "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
@@ -23,57 +24,49 @@ 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
-            conf = float(box.conf[0])  # Confidence score
-            
-            # 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}
+            label = CLASS_NAMES[class_id]  # Get class label
+            confidence = box.conf[0].item()  # Confidence score
 
-    # 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 **BELOW** text elements
+            cv2.rectangle(annotated_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
 
-        # Draw bounding box
-        cv2.rectangle(annotated_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
+            # Set text properties
+            font = cv2.FONT_HERSHEY_SIMPLEX
+            font_scale = 0.7  # Increased font size and thickness for better readability
+            font_thickness = 1.3
+            label_text = f"{label} ({confidence:.2f})"
 
-        # 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
+            # Get text size for proper background padding
+            text_size = cv2.getTextSize(label_text, 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)
+            # 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
-        )
+            # Draw **filled rectangle background** for the text (above bounding box)
+            cv2.rectangle(
+                annotated_image, 
+                (text_x, text_y - text_size[1] - 5), 
+                (text_x + text_size[0] + 5, text_y + 5), 
+                (0, 255, 0),  # Background color (Green)
+                -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
-        )
+            # Draw the **text label above the rectangle**
+            cv2.putText(
+                annotated_image, 
+                label_text, 
+                (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