Update app.py

app.py

@@ -59,7 +59,7 @@ def _text_size(draw, text, font):
     return bbox[2] - bbox[0], bbox[3] - bbox[1]
 
 
-# ─── Visual polygon helpers ───────────────────────────────────────────────────
+# ─── Mask helpers ─────────────────────────────────────────────────────────────
 
 def _polygon_to_mask(pts_xy, h, w):
     """Rasterise raw polygon → binary uint8 mask. BACKEND / measurements only."""
@@ -69,19 +69,90 @@ def _polygon_to_mask(pts_xy, h, w):
     return mask
 
 
-def _rotated_box_from_mask(mask_np):
+def _refine_mask_grabcut(img_bgr, coarse_mask):
     """
-    Compute tight rotated bounding box from backend mask.
-    Returns the 4 corner points as int32 array (4, 1, 2) — ready for
-    cv2.polylines / fillPoly.
+    Refine a coarse binary mask to pixel-perfect grain boundary using GrabCut.
+    img_bgr     : full BGR image
+    coarse_mask : uint8 binary mask (0/1), same size as img_bgr
+    Returns     : refined binary uint8 mask (0/1)
     """
-    pts_y, pts_x = np.where(mask_np == 1)
-    if len(pts_x) < 5:
-        return None, None
-    pts = np.column_stack([pts_x.astype(np.float32), pts_y.astype(np.float32)])
-    rect = cv2.minAreaRect(pts)
-    box  = cv2.boxPoints(rect)          # 4 corners (float)
-    return box.astype(np.int32).reshape(-1, 1, 2), rect
+    ys, xs = np.where(coarse_mask == 1)
+    if len(xs) < 5:
+        return coarse_mask
+
+    # Tight crop with small padding so GrabCut has background context
+    x1, y1 = max(0, int(xs.min()) - 6), max(0, int(ys.min()) - 6)
+    x2, y2 = min(img_bgr.shape[1], int(xs.max()) + 6), min(img_bgr.shape[0], int(ys.max()) + 6)
+    crop    = img_bgr[y1:y2, x1:x2]
+    ch, cw  = crop.shape[:2]
+    if ch < 8 or cw < 8:
+        return coarse_mask
+
+    # Build GrabCut init mask from coarse mask crop
+    gc_mask  = np.full((ch, cw), cv2.GC_BGD, dtype=np.uint8)
+    local_fg = coarse_mask[y1:y2, x1:x2]
+
+    # Erode to get definite FG core, dilate to get probable FG ring
+    k_sm = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
+    k_lg = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (11, 11))
+    def_fg   = cv2.erode(local_fg,  k_sm, iterations=2)
+    prob_fg  = cv2.dilate(local_fg, k_lg, iterations=2)
+
+    gc_mask[prob_fg == 1] = cv2.GC_PR_FGD
+    gc_mask[def_fg  == 1] = cv2.GC_FGD
+    # Border strip = definite background
+    gc_mask[:3, :]  = cv2.GC_BGD
+    gc_mask[-3:, :] = cv2.GC_BGD
+    gc_mask[:, :3]  = cv2.GC_BGD
+    gc_mask[:, -3:] = cv2.GC_BGD
+
+    try:
+        bgd_model = np.zeros((1, 65), np.float64)
+        fgd_model = np.zeros((1, 65), np.float64)
+        cv2.grabCut(crop, gc_mask, None, bgd_model, fgd_model, 4, cv2.GC_INIT_WITH_MASK)
+        refined_local = np.where((gc_mask == cv2.GC_FGD) | (gc_mask == cv2.GC_PR_FGD), 1, 0).astype(np.uint8)
+    except Exception:
+        return coarse_mask
+
+    # Clean up with morphology: close small holes, smooth jagged edges
+    k_cl = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
+    refined_local = cv2.morphologyEx(refined_local, cv2.MORPH_CLOSE, k_cl, iterations=2)
+    refined_local = cv2.morphologyEx(refined_local, cv2.MORPH_OPEN,  k_cl, iterations=1)
+
+    # Put refined crop back into full-size mask
+    refined_full = np.zeros_like(coarse_mask)
+    refined_full[y1:y2, x1:x2] = refined_local
+    return refined_full
+
+
+def _mask_to_smooth_contour(mask_np):
+    """
+    Extract the outer contour of a binary mask and smooth it with
+    spline-like resampling → returns int32 array (N,1,2) for cv2 drawing.
+    """
+    contours, _ = cv2.findContours(mask_np, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
+    if not contours:
+        return None
+    cnt = max(contours, key=cv2.contourArea).astype(np.float32).reshape(-1, 2)
+    if len(cnt) < 6:
+        return cnt.astype(np.int32).reshape(-1, 1, 2)
+
+    # Resample to ~120 evenly-spaced points for a smooth outline
+    n_target = min(120, max(40, len(cnt)))
+    indices  = np.linspace(0, len(cnt) - 1, n_target).astype(int)
+    sampled  = cnt[indices]
+
+    # Circular Gaussian smooth
+    window   = 9
+    half     = window // 2
+    padded   = np.vstack([sampled[-half:], sampled, sampled[:half]])
+    kernel   = cv2.getGaussianKernel(window, 0).flatten()
+    kernel  /= kernel.sum()
+    smoothed = np.zeros_like(sampled)
+    for i in range(len(sampled)):
+        smoothed[i] = (padded[i:i + window] * kernel[:, None]).sum(axis=0)
+
+    return smoothed.astype(np.int32).reshape(-1, 1, 2)
 
 
 # ─────────────────────────────────────────────────────────────────────────────
@@ -93,6 +164,7 @@ def _rotated_box_from_mask(mask_np):
 # ──────────────────────────────────��──────────────────────────────────────────
 def run_segmentation(img_np):
     h, w    = img_np.shape[:2]
+    img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
     results = model(img_np, imgsz=1280, conf=0.25)[0]
 
     annotated   = img_np.copy()
@@ -114,11 +186,12 @@ def run_segmentation(img_np):
             color    = CLASS_COLORS.get(cls_id, (200, 200, 200))
             counts[cls_name] += 1
 
-            # Backend mask: raw polygon fill (measurements only)
+            # Backend mask: raw polygon fill (used for measurements — never changed)
             mask_np = _polygon_to_mask(poly_xy, h, w)
 
-            # Visual: tight rotated bounding box from the backend mask
-            vis_box, rect = _rotated_box_from_mask(mask_np)
+            # Visual mask: GrabCut-refined → hugs actual grain pixels perfectly
+            vis_mask    = _refine_mask_grabcut(img_bgr, mask_np)
+            vis_contour = _mask_to_smooth_contour(vis_mask)
 
             # Bounding box from backend mask for zoom crop
             ys, xs = np.where(mask_np == 1)
@@ -128,24 +201,25 @@ def run_segmentation(img_np):
                 all_x2 = max(all_x2, int(xs.max()))
                 all_y2 = max(all_y2, int(ys.max()))
 
-            if vis_box is not None:
-                cv2.fillPoly(overlay, [vis_box], color)
+            # Fill overlay using the refined visual mask directly (pixel-perfect fill)
+            overlay[vis_mask == 1] = color
 
             grain_boxes.append({
-                "cls_id":   cls_id,
-                "cls_name": cls_name,
-                "mask_np":  mask_np,    # backend only
-                "vis_box":  vis_box,    # visual rotated bbox
+                "cls_id":     cls_id,
+                "cls_name":   cls_name,
+                "mask_np":    mask_np,      # backend only — measurements
+                "vis_mask":   vis_mask,     # refined visual mask
+                "vis_contour": vis_contour, # smooth contour for outline
             })
 
     # Blend fill
     annotated = cv2.addWeighted(annotated, 0.72, overlay, 0.28, 0)
 
-    # Draw crisp rotated bbox outlines (no fill — just the border)
+    # Draw smooth anti-aliased contour outlines over the blend
     for g in grain_boxes:
-        if g["vis_box"] is not None:
+        if g["vis_contour"] is not None:
             cv2.polylines(
-                annotated, [g["vis_box"]],
+                annotated, [g["vis_contour"]],
                 isClosed=True, color=CLASS_COLORS[g["cls_id"]], thickness=2,
                 lineType=cv2.LINE_AA,
             )