Upload app.py

app.py

@@ -86,26 +86,9 @@ def draw_boxes(img_rgb, boxes, labels, color=(0, 215, 255)):
     for box, label in zip(boxes, labels):
         x1, y1, x2, y2 = map(int, box)
         
-        # Fancy stylized corners for the bounding box
-        length = int(min(x2 - x1, y2 - y1) * 0.2)
-        thickness = 3
+        # Solid continuous bounding box
+        cv2.rectangle(out, (x1, y1), (x2, y2), color, 2)
         
-        # Draw faint full rectangle
-        cv2.rectangle(out, (x1, y1), (x2, y2), color, 1)
-        
-        # Top-left corner
-        cv2.line(out, (x1, y1), (x1 + length, y1), color, thickness, cv2.LINE_AA)
-        cv2.line(out, (x1, y1), (x1, y1 + length), color, thickness, cv2.LINE_AA)
-        # Top-right corner
-        cv2.line(out, (x2, y1), (x2 - length, y1), color, thickness, cv2.LINE_AA)
-        cv2.line(out, (x2, y1), (x2, y1 + length), color, thickness, cv2.LINE_AA)
-        # Bottom-left corner
-        cv2.line(out, (x1, y2), (x1 + length, y2), color, thickness, cv2.LINE_AA)
-        cv2.line(out, (x1, y2), (x1, y2 - length), color, thickness, cv2.LINE_AA)
-        # Bottom-right corner
-        cv2.line(out, (x2, y2), (x2 - length, y2), color, thickness, cv2.LINE_AA)
-        cv2.line(out, (x2, y2), (x2, y2 - length), color, thickness, cv2.LINE_AA)
-
         # Draw a solid background for the text label for better readability
         (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
         cv2.rectangle(out, (x1, max(y1 - 25, 0)), (x1 + w + 10, max(y1 - 25, 0) + h + 10), color, -1)
@@ -127,23 +110,26 @@ def run_yolo_generic(img_rgb, model_path, target_classes, color):
 
     result = results[0]
     h, w = img_rgb.shape[:2]
-    combined_mask = np.zeros((h, w), dtype=bool)
+    combined_mask = np.zeros((h, w), dtype=np.uint8)
     boxes, labels = [], []
 
     if result.masks is not None:
-        for mask, box, cls, conf in zip(
-            result.masks.data, result.boxes.xyxy,
+        for pts, box, cls, conf in zip(
+            result.masks.xy, result.boxes.xyxy,
             result.boxes.cls, result.boxes.conf
         ):
             if int(cls) not in target_classes:
                 continue
-            mask_np = mask.cpu().numpy()
-            mask_r  = cv2.resize(mask_np, (w, h), interpolation=cv2.INTER_NEAREST)
-            combined_mask |= mask_r > 0.5
+            
+            # Use polygon points perfectly mapped to original image dimensions
+            if len(pts) > 0:
+                cv2.fillPoly(combined_mask, [np.int32(pts)], 1)
+                
             boxes.append(box.cpu().tolist())
             labels.append(f"glass {conf:.2f}")
 
-    out = apply_mask_overlay(img_rgb, combined_mask, color=color)
+    combined_mask_bool = combined_mask > 0
+    out = apply_mask_overlay(img_rgb, combined_mask_bool, color=color)
     out = draw_boxes(out, boxes, labels, color=color)
     bw_mask = (combined_mask * 255).astype(np.uint8)
     return out, bw_mask, f"Found: {len(boxes)} | Inference Time: {elapsed:.2f}s"