Update app.py

app.py

@@ -58,17 +58,60 @@ def _text_size(draw, text, font):
     bbox = draw.textbbox((0, 0), text, font=font)
     return bbox[2] - bbox[0], bbox[3] - bbox[1]
 
+
+# ─── Visual-only helpers ───────────────────────────────────────────────────────
+def _make_visual_mask(mask_raw_float, w, h):
+    """
+    Upsample the raw float mask (0..1) with bilinear interpolation and apply
+    morphological smoothing.  Result is a uint8 binary mask for DISPLAY ONLY
+    — backend mask_np (INTER_NEAREST) is kept separate for measurements.
+    """
+    # Bilinear upscale → much better edge alignment than NEAREST
+    vis = cv2.resize(mask_raw_float, (w, h), interpolation=cv2.INTER_LINEAR)
+    vis = (vis > 0.45).astype(np.uint8)
+
+    # Elliptical close then open: fills jagged notches, shaves sharp spurs
+    k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
+    vis = cv2.morphologyEx(vis, cv2.MORPH_CLOSE, k)
+    vis = cv2.morphologyEx(vis, cv2.MORPH_OPEN,  k)
+    return vis
+
+
+def _smooth_contour(contour, window=9):
+    """
+    Circular sliding-window average over contour points.
+    Smooths wiggly edges while preserving overall shape / corners.
+    Returns an int32 contour array compatible with cv2.polylines.
+    """
+    pts = contour[:, 0, :].astype(np.float32)
+    n   = len(pts)
+    if n < window * 2 + 1:
+        return contour.astype(np.int32)
+
+    half    = window // 2
+    padded  = np.vstack([pts[-half:], pts, pts[:half]])
+    weights = np.hanning(window).astype(np.float32)
+    weights /= weights.sum()
+
+    smoothed = np.zeros_like(pts)
+    for i in range(n):
+        smoothed[i] = (padded[i : i + window] * weights[:, None]).sum(axis=0)
+
+    return smoothed.astype(np.int32).reshape(-1, 1, 2)
+
+
 # ─────────────────────────────────────────────────────────────────────────────
 # STEP 1 — Fast segmentation + classification output image (no text labels)
 #           + auto-crop/zoom to grain bounding box
 # ─────────────────────────────────────────────────────────────────────────────
 def run_segmentation(img_np):
     """
-    Run YOLO, draw colored masks + contours (NO text labels on image).
+    Run YOLO, draw colored masks + smooth contours (NO text labels on image).
     Returns:
       annotated    : full-size RGB annotated image (numpy)
       zoomed_pil   : PIL image cropped/zoomed to grain region
       grain_boxes  : list of dicts {cls_id, cls_name, mask_np, bbox}
+                     mask_np is INTER_NEAREST — used only for measurements.
       counts       : {"Full": int, "Broken": int}
     """
     h, w    = img_np.shape[:2]
@@ -79,7 +122,7 @@ def run_segmentation(img_np):
     counts     = {"Full": 0, "Broken": 0}
     grain_boxes = []
 
-    # Bounding box of all grains combined (for zoom)
+    # Bounding box of all grains combined (for zoom) — uses backend mask
     all_x1, all_y1, all_x2, all_y2 = w, h, 0, 0
 
     if results.masks is not None:
@@ -89,10 +132,16 @@ def run_segmentation(img_np):
             color    = CLASS_COLORS.get(cls_id, (200, 200, 200))
             counts[cls_name] += 1
 
-            mask_np = mask_tensor.cpu().numpy().astype(np.uint8)
-            mask_np = cv2.resize(mask_np, (w, h), interpolation=cv2.INTER_NEAREST)
+            # ── Backend mask (INTER_NEAREST) — used for measurements only ──
+            mask_raw = mask_tensor.cpu().numpy()           # float32, 0..1
+            mask_np  = cv2.resize(
+                mask_raw, (w, h), interpolation=cv2.INTER_NEAREST
+            ).astype(np.uint8)
+
+            # ── Visual mask (bilinear + morphological smooth) ──────────────
+            mask_vis = _make_visual_mask(mask_raw, w, h)
 
-            # Update combined grain bounding box
+            # Update combined grain bounding box from backend mask
             ys, xs = np.where(mask_np == 1)
             if len(xs) > 0:
                 all_x1 = min(all_x1, int(xs.min()))
@@ -100,25 +149,41 @@ def run_segmentation(img_np):
                 all_x2 = max(all_x2, int(xs.max()))
                 all_y2 = max(all_y2, int(ys.max()))
 
-            overlay[mask_np == 1] = color
-            cnts, _ = cv2.findContours(mask_np, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-            cv2.drawContours(annotated, cnts, -1, color, 2)
+            # Fill overlay using visual mask (smooth region)
+            overlay[mask_vis == 1] = color
 
             grain_boxes.append({
-                "cls_id":   cls_id,
-                "cls_name": cls_name,
-                "mask_np":  mask_np,
+                "cls_id":    cls_id,
+                "cls_name":  cls_name,
+                "mask_np":   mask_np,    # ← backend only, never touched for visuals
+                "mask_vis":  mask_vis,   # ← visual only, never used for measurements
             })
 
-    # Blend overlays
+    # Blend fill (overlay) into annotated
     annotated = cv2.addWeighted(annotated, 0.72, overlay, 0.28, 0)
 
-    # Redraw contours sharp on top
+    # ── Draw smooth, anti-aliased contours on top ──────────────────────────
     for g in grain_boxes:
-        cnts, _ = cv2.findContours(g["mask_np"], cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-        cv2.drawContours(annotated, cnts, -1, CLASS_COLORS[g["cls_id"]], 2)
+        color = CLASS_COLORS[g["cls_id"]]
+        cnts, _ = cv2.findContours(
+            g["mask_vis"], cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
+        )
+        smooth_cnts = []
+        for c in cnts:
+            if len(c) >= 10:
+                smooth_cnts.append(_smooth_contour(c, window=9))
+            elif len(c) >= 4:
+                smooth_cnts.append(c.astype(np.int32))
+
+        if smooth_cnts:
+            # polylines with LINE_AA gives sub-pixel anti-aliased edges
+            cv2.polylines(
+                annotated, smooth_cnts,
+                isClosed=True, color=color, thickness=2,
+                lineType=cv2.LINE_AA,
+            )
 
-    # ── Zoom to grain region with padding ─────────────────────────────────────
+    # ── Zoom to grain region with padding ─────────────────────────────────
     if all_x2 > all_x1 and all_y2 > all_y1:
         pad   = max(30, int(max(all_x2 - all_x1, all_y2 - all_y1) * 0.08))
         cx1   = max(0, all_x1 - pad)
@@ -142,12 +207,13 @@ def measure_grains_from_boxes(grain_boxes, img_shape, paper_dims):
     h_px = long axis, w_px = short axis.
     Converts to mm if paper_dims is provided.
     Returns list of measurement dicts.
+    NOTE: uses mask_np (INTER_NEAREST) — measurements are unaffected by visual smoothing.
     """
     paper_px = (paper_dims[0] + paper_dims[1]) / 2.0 if paper_dims else None
     measurements = []
 
     for idx, g in enumerate(grain_boxes):
-        mask_np = g["mask_np"]
+        mask_np = g["mask_np"]          # backend mask only
         pts_y, pts_x = np.where(mask_np == 1)
         if len(pts_x) < 5:
             continue