Update app.py

app.py

@@ -15,7 +15,7 @@ CLASS_COLORS = {0: (34, 197, 94), 1: (239, 68, 68)}   # green, red
 SAMPLE_PATHS = ["image1.jpg", "image2.jpg"]
 
 # ─── Paper reference ──────────────────────────────────────────────────────────
-PAPER_REAL_MM = 40.0   # white 4×4 cm square = 40 mm per side
+PAPER_REAL_MM = 40.0   # white 4x4 cm square = 40 mm per side
 
 def detect_paper_pixels(img_np):
     gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
@@ -34,7 +34,7 @@ def detect_paper_pixels(img_np):
         if 0.5 < (w / max(h, 1)) < 2.0 and area > best_area:
             best_area = area
             best = (h, w)
-    return best  # (h_px, w_px) or None
+    return best
 
 def px_to_mm(px, paper_px_dim):
     if not paper_px_dim:
@@ -59,36 +59,25 @@ def _text_size(draw, text, 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
+# ─── Visual polygon helpers ────────────────────────────────────────────────────
 
-
-def _smooth_contour(contour, window=9):
+def _smooth_polygon(pts_xy, window=11):
     """
-    Circular sliding-window average over contour points.
-    Smooths wiggly edges while preserving overall shape / corners.
-    Returns an int32 contour array compatible with cv2.polylines.
+    Hanning-weighted circular sliding-window average on polygon vertices.
+    pts_xy  : numpy (N, 2) float — already in original image coordinates
+              from results.masks.xy, so NO resize drift whatsoever.
+    Returns : int32 array (N, 1, 2) for cv2.polylines / fillPoly.
+    VISUAL ONLY — backend mask uses raw polygon via _polygon_to_mask().
     """
-    pts = contour[:, 0, :].astype(np.float32)
+    pts = pts_xy.astype(np.float32)
     n   = len(pts)
-    if n < window * 2 + 1:
-        return contour.astype(np.int32)
+    if n < 6:
+        return pts.astype(np.int32).reshape(-1, 1, 2)
+
+    # Window must be odd and fit inside the polygon
+    window = min(window | 1, (n - 1) | 1)
+    half   = window // 2
 
-    half    = window // 2
     padded  = np.vstack([pts[-half:], pts, pts[:half]])
     weights = np.hanning(window).astype(np.float32)
     weights /= weights.sum()
@@ -100,48 +89,58 @@ def _smooth_contour(contour, window=9):
     return smoothed.astype(np.int32).reshape(-1, 1, 2)
 
 
+def _polygon_to_mask(pts_xy, h, w):
+    """
+    Rasterise raw polygon to binary uint8 mask.
+    Used ONLY for backend measurements — never for visuals.
+    """
+    mask = np.zeros((h, w), dtype=np.uint8)
+    if len(pts_xy) >= 3:
+        cv2.fillPoly(mask, [pts_xy.astype(np.int32)], 1)
+    return mask
+
+
 # ─────────────────────────────────────────────────────────────────────────────
-# STEP 1 — Fast segmentation + classification output image (no text labels)
-#           + auto-crop/zoom to grain bounding box
+# STEP 1 — Segmentation + visual output
+#
+# ROOT-CAUSE FIX FOR MISALIGNMENT:
+#   Previously used results.masks.data (low-res tensor) + cv2.resize
+#   which introduces sub-pixel drift at every grain boundary.
+#
+#   Now uses results.masks.xy — ultralytics already maps each polygon
+#   to ORIGINAL image pixel coordinates, so alignment is exact.
+#   No resize, no drift, no displacement.
 # ─────────────────────────────────────────────────────────────────────────────
 def run_segmentation(img_np):
-    """
-    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]
     results = model(img_np, imgsz=1280, conf=0.25)[0]
 
-    annotated  = img_np.copy()
-    overlay    = img_np.copy()
-    counts     = {"Full": 0, "Broken": 0}
+    annotated   = img_np.copy()
+    overlay     = img_np.copy()
+    counts      = {"Full": 0, "Broken": 0}
     grain_boxes = []
 
-    # 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:
-        for mask_tensor, box in zip(results.masks.data, results.boxes):
+        xy_list = results.masks.xy          # list of (N_i, 2) arrays, original coords
+
+        for poly_xy, box in zip(xy_list, results.boxes):
+            if len(poly_xy) < 3:
+                continue
+
             cls_id   = int(box.cls[0])
             cls_name = CLASS_NAMES.get(cls_id, "?")
             color    = CLASS_COLORS.get(cls_id, (200, 200, 200))
             counts[cls_name] += 1
 
-            # ── 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)
+            # Backend mask: raw polygon fill — for measurements only
+            mask_np = _polygon_to_mask(poly_xy, h, w)
 
-            # ── Visual mask (bilinear + morphological smooth) ──────────────
-            mask_vis = _make_visual_mask(mask_raw, w, h)
+            # Visual polygon: Hanning-smoothed — for display only
+            vis_poly = _smooth_polygon(poly_xy, window=11)
 
-            # Update combined grain bounding box from backend mask
+            # Update zoom bounding box from backend mask
             ys, xs = np.where(mask_np == 1)
             if len(xs) > 0:
                 all_x1 = min(all_x1, int(xs.min()))
@@ -149,49 +148,36 @@ def run_segmentation(img_np):
                 all_x2 = max(all_x2, int(xs.max()))
                 all_y2 = max(all_y2, int(ys.max()))
 
-            # Fill overlay using visual mask (smooth region)
-            overlay[mask_vis == 1] = color
+            # Fill overlay with smooth visual polygon
+            cv2.fillPoly(overlay, [vis_poly], color)
 
             grain_boxes.append({
-                "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
+                "cls_id":   cls_id,
+                "cls_name": cls_name,
+                "mask_np":  mask_np,     # backend only
+                "vis_poly": vis_poly,    # visual only
             })
 
-    # Blend fill (overlay) into annotated
+    # Blend fill
     annotated = cv2.addWeighted(annotated, 0.72, overlay, 0.28, 0)
 
-    # ── Draw smooth, anti-aliased contours on top ──────────────────────────
+    # Draw smooth anti-aliased outlines on top
     for g in grain_boxes:
         color = CLASS_COLORS[g["cls_id"]]
-        cnts, _ = cv2.findContours(
-            g["mask_vis"], cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
+        cv2.polylines(
+            annotated, [g["vis_poly"]],
+            isClosed=True, color=color, thickness=2,
+            lineType=cv2.LINE_AA,
         )
-        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
     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)
-        cy1   = max(0, all_y1 - pad)
-        cx2   = min(w, all_x2 + pad)
-        cy2   = min(h, all_y2 + pad)
-        crop  = annotated[cy1:cy2, cx1:cx2]
-        zoomed_pil = Image.fromarray(crop)
+        pad  = max(30, int(max(all_x2 - all_x1, all_y2 - all_y1) * 0.08))
+        cx1  = max(0, all_x1 - pad)
+        cy1  = max(0, all_y1 - pad)
+        cx2  = min(w, all_x2 + pad)
+        cy2  = min(h, all_y2 + pad)
+        zoomed_pil = Image.fromarray(annotated[cy1:cy2, cx1:cx2])
     else:
         zoomed_pil = Image.fromarray(annotated)
 
@@ -199,21 +185,14 @@ def run_segmentation(img_np):
 
 
 # ─────────────────────────────────────────────────────────────────────────────
-# STEP 2 — Measure grain height & width from YOLO masks
+# STEP 2 — Measure grains (unchanged — uses backend mask_np only)
 # ─────────────────────────────────────────────────────────────────────────────
 def measure_grains_from_boxes(grain_boxes, img_shape, paper_dims):
-    """
-    For each grain mask, compute oriented bounding box (minAreaRect).
-    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"]          # backend mask only
+        mask_np = g["mask_np"]
         pts_y, pts_x = np.where(mask_np == 1)
         if len(pts_x) < 5:
             continue
@@ -222,10 +201,10 @@ def measure_grains_from_boxes(grain_boxes, img_shape, paper_dims):
         rect = cv2.minAreaRect(pts)
         (cx, cy), (rw, rh), _ = rect
 
-        h_px = float(max(rw, rh))
-        w_px = float(min(rw, rh))
-        h_mm = px_to_mm(h_px, paper_px)
-        w_mm = px_to_mm(w_px, paper_px)
+        h_px     = float(max(rw, rh))
+        w_px     = float(min(rw, rh))
+        h_mm     = px_to_mm(h_px, paper_px)
+        w_mm     = px_to_mm(w_px, paper_px)
         area_mm2 = (h_mm * w_mm) if (h_mm and w_mm) else None
 
         measurements.append({
@@ -244,7 +223,7 @@ def measure_grains_from_boxes(grain_boxes, img_shape, paper_dims):
 
 
 # ─────────────────────────────────────────────────────────────────────────────
-# STEP 2b — Build measurement DataFrames (same style as previous.py)
+# STEP 2b — Build DataFrames
 # ─────────────────────────────────────────────────────────────────────────────
 def build_table_data(measurements, paper_px, counts):
     has_mm = paper_px is not None
@@ -252,28 +231,27 @@ def build_table_data(measurements, paper_px, counts):
 
     rows = []
     for g in measurements:
-        h_val    = round(g["h_mm"], 2)  if (has_mm and g["h_mm"])    else round(g["h_px"], 1)
-        w_val    = round(g["w_mm"], 2)  if (has_mm and g["w_mm"])    else round(g["w_px"], 1)
+        h_val    = round(g["h_mm"],    2) if (has_mm and g["h_mm"])    else round(g["h_px"], 1)
+        w_val    = round(g["w_mm"],    2) if (has_mm and g["w_mm"])    else round(g["w_px"], 1)
         area_val = round(g["area_mm2"], 2) if g["area_mm2"] else None
         rows.append({
-            "#":                        g["label"],
-            "Class":                    g["cls_name"],
-            f"Height ({unit})":         h_val,
-            f"Width ({unit})":          w_val,
-            f"Area (mm²)" if has_mm else "Area": area_val,
+            "#":                                 g["label"],
+            "Class":                             g["cls_name"],
+            f"Height ({unit})":                  h_val,
+            f"Width ({unit})":                   w_val,
+            "Area (mm\u00b2)" if has_mm else "Area": area_val,
         })
     grain_df = pd.DataFrame(rows)
 
-    # Summary
-    h_key = "h_mm" if has_mm else "h_px"
-    w_key = "w_mm" if has_mm else "w_px"
+    h_key   = "h_mm" if has_mm else "h_px"
+    w_key   = "w_mm" if has_mm else "w_px"
     heights = [(g["label"], g[h_key]) for g in measurements if g.get(h_key)]
     widths  = [(g["label"], g[w_key]) for g in measurements if g.get(w_key)]
 
-    max_h = max(heights, key=lambda x: x[1]) if heights else (0, 0)
-    min_h = min(heights, key=lambda x: x[1]) if heights else (0, 0)
-    max_w = max(widths,  key=lambda x: x[1]) if widths  else (0, 0)
-    min_w = min(widths,  key=lambda x: x[1]) if widths  else (0, 0)
+    max_h    = max(heights, key=lambda x: x[1]) if heights else (0, 0)
+    min_h    = min(heights, key=lambda x: x[1]) if heights else (0, 0)
+    max_w    = max(widths,  key=lambda x: x[1]) if widths  else (0, 0)
+    min_w    = min(widths,  key=lambda x: x[1]) if widths  else (0, 0)
     interval = (max_h[1] - min_h[1]) / 10.0 if (heights and max_h[1] != min_h[1]) else 0.0
 
     n_full   = counts.get("Full",   0)
@@ -281,64 +259,49 @@ def build_table_data(measurements, paper_px, counts):
     total    = n_full + n_broken
 
     summary_rows = [
-        {"Metric": "Total Grains",                          "Value": str(total)},
-        {"Metric": "🟢 Full Grains",                        "Value": str(n_full)},
-        {"Metric": "🔴 Broken Grains",                      "Value": str(n_broken)},
-        {"Metric": "Paper Reference",                       "Value": f"✅ Found ({unit} mode)" if has_mm else "❌ Not found (px only)"},
-        {"Metric": f"Max Height (Grain #{max_h[0]})",       "Value": f"{max_h[1]:.2f} {unit}"},
-        {"Metric": f"Min Height (Grain #{min_h[0]})",       "Value": f"{min_h[1]:.2f} {unit}"},
-        {"Metric": f"Max Width  (Grain #{max_w[0]})",       "Value": f"{max_w[1]:.2f} {unit}"},
-        {"Metric": f"Min Width  (Grain #{min_w[0]})",       "Value": f"{min_w[1]:.2f} {unit}"},
-        {"Metric": f"Mean Height",                          "Value": f"{np.mean([v for _, v in heights]):.2f} {unit}" if heights else "—"},
-        {"Metric": f"Mean Width",                           "Value": f"{np.mean([v for _, v in widths]):.2f} {unit}"  if widths  else "—"},
-        {"Metric": "Bin Interval (max−min)/10",             "Value": f"{interval:.3f} {unit}"},
+        {"Metric": "Total Grains",                    "Value": str(total)},
+        {"Metric": "🟢 Full Grains",                  "Value": str(n_full)},
+        {"Metric": "🔴 Broken Grains",                "Value": str(n_broken)},
+        {"Metric": "Paper Reference",                 "Value": f"✅ Found ({unit} mode)" if has_mm else "❌ Not found (px only)"},
+        {"Metric": f"Max Height (Grain #{max_h[0]})", "Value": f"{max_h[1]:.2f} {unit}"},
+        {"Metric": f"Min Height (Grain #{min_h[0]})", "Value": f"{min_h[1]:.2f} {unit}"},
+        {"Metric": f"Max Width  (Grain #{max_w[0]})", "Value": f"{max_w[1]:.2f} {unit}"},
+        {"Metric": f"Min Width  (Grain #{min_w[0]})", "Value": f"{min_w[1]:.2f} {unit}"},
+        {"Metric": "Mean Height",                     "Value": f"{np.mean([v for _, v in heights]):.2f} {unit}" if heights else "—"},
+        {"Metric": "Mean Width",                      "Value": f"{np.mean([v for _, v in widths]):.2f} {unit}"  if widths  else "—"},
+        {"Metric": "Bin Interval (max-min)/10",       "Value": f"{interval:.3f} {unit}"},
     ]
     summary_df = pd.DataFrame(summary_rows)
     return grain_df, summary_df
 
 
 # ─────────────────────────────────────────────────────────────────────────────
-# GRADIO — TWO-STAGE predict
-# Stage 1: fast — returns segmentation image + count summary (no tables yet)
-# Stage 2: slower — returns measurement tables
+# GRADIO — two-stage predict
 # ─────────────────────────────────────────────────────────────────────────────
 def predict_stage1(image: Image.Image):
-    """Returns zoomed segmentation image + summary + placeholder table message."""
     if image is None:
         return None, "", "", None, None
-
     img_np = np.array(image.convert("RGB"))
-
-    # Run segmentation (fast)
     _, zoomed_pil, grain_boxes, counts = run_segmentation(img_np)
-
-    total   = counts["Full"] + counts["Broken"]
-    summary = f"✅  {total} grains detected  ·  🟢 Full: {counts['Full']}  ·  🔴 Broken: {counts['Broken']}"
-
-    count_md = f"""| | Count |
-|---|---|
-| 🌾 Total Grains | **{total}** |
-| 🟢 Full Grains | **{counts['Full']}** |
-| 🔴 Broken Grains | **{counts['Broken']}** |"""
-
-    # Return image + summary immediately; tables = loading placeholder
+    total    = counts["Full"] + counts["Broken"]
+    summary  = f"✅  {total} grains detected  ·  🟢 Full: {counts['Full']}  ·  🔴 Broken: {counts['Broken']}"
+    count_md = (
+        f"| | Count |\n|---|---|\n"
+        f"| 🌾 Total Grains | **{total}** |\n"
+        f"| 🟢 Full Grains | **{counts['Full']}** |\n"
+        f"| 🔴 Broken Grains | **{counts['Broken']}** |"
+    )
     loading_df = pd.DataFrame([{"Status": "⏳  Calculating height & width of all grains..."}])
     return zoomed_pil, summary, count_md, loading_df, loading_df
 
 
 def predict_stage2(image: Image.Image):
-    """Full pipeline — returns everything including measurement tables."""
     if image is None:
         return None, "", "", None, None
-
     img_np = np.array(image.convert("RGB"))
 
-    # Run segmentation + paper detection in parallel
-    def _seg():
-        return run_segmentation(img_np)
-
-    def _paper():
-        return detect_paper_pixels(img_np)
+    def _seg():   return run_segmentation(img_np)
+    def _paper(): return detect_paper_pixels(img_np)
 
     with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
         fut_seg   = pool.submit(_seg)
@@ -346,26 +309,24 @@ def predict_stage2(image: Image.Image):
         _, zoomed_pil, grain_boxes, counts = fut_seg.result()
         paper_dims = fut_paper.result()
 
-    # Measure grains
     measurements, paper_px = measure_grains_from_boxes(grain_boxes, img_np.shape, paper_dims)
-
-    total   = counts["Full"] + counts["Broken"]
-    summary = (f"✅  {total} grains detected  ·  🟢 Full: {counts['Full']}  ·  🔴 Broken: {counts['Broken']}"
-               + (f"  ·  📏 Paper found — measurements in mm" if paper_px else "  ·  ⚠️ No paper — measurements in px"))
-
-    count_md = f"""| | Count |
-|---|---|
-| 🌾 Total Grains | **{total}** |
-| 🟢 Full Grains | **{counts['Full']}** |
-| 🔴 Broken Grains | **{counts['Broken']}** |"""
-
+    total    = counts["Full"] + counts["Broken"]
+    summary  = (
+        f"✅  {total} grains detected  ·  🟢 Full: {counts['Full']}  ·  🔴 Broken: {counts['Broken']}"
+        + (f"  ·  📏 Paper found — measurements in mm" if paper_px else "  ·  ⚠️ No paper — measurements in px")
+    )
+    count_md = (
+        f"| | Count |\n|---|---|\n"
+        f"| 🌾 Total Grains | **{total}** |\n"
+        f"| 🟢 Full Grains | **{counts['Full']}** |\n"
+        f"| 🔴 Broken Grains | **{counts['Broken']}** |"
+    )
     grain_df, summary_df = build_table_data(measurements, paper_px, counts)
-
     return zoomed_pil, summary, count_md, grain_df, summary_df
 
 
 # ─────────────────────────────────────────────────────────────────────────────
-# UI  (theme + styling identical to previous.py)
+# UI
 # ─────────────────────────────────────────────────────────────────────────────
 THEME = gr.themes.Soft(
     primary_hue="violet",
@@ -377,7 +338,6 @@ THEME = gr.themes.Soft(
 CSS = """
 #run-btn { margin-top: 6px; }
 #status-box textarea { font-size: 0.92rem; }
-.gr-dataframe table  { font-size: 0.88rem; }
 #count-box { font-size: 0.95rem; }
 """
 
@@ -397,33 +357,25 @@ with gr.Blocks(theme=THEME, title="GrainVision", css=CSS) as demo:
     """)
 
     with gr.Row(equal_height=False):
-
-        # ── LEFT: Input ───────────────────────────────────────────────────────
         with gr.Column(scale=1):
             inp_image = gr.Image(type="pil", label="Upload Rice Image", height=280)
-
-            run_btn = gr.Button("🔍  Analyse Grains",
-                                variant="primary", size="lg", elem_id="run-btn")
-
+            run_btn   = gr.Button("🔍  Analyse Grains",
+                                  variant="primary", size="lg", elem_id="run-btn")
             gr.Markdown("_Upload an image then press **Analyse**. Segmentation appears first, measurements follow._")
-
             status_box = gr.Textbox(
                 label="Status", value="", interactive=False,
                 visible=True, max_lines=3, elem_id="status-box",
             )
-
             gr.Markdown("### Example Images  _(click to load)_")
             gr.Examples(
                 examples=[[p] for p in SAMPLE_PATHS],
                 inputs=inp_image, label="", examples_per_page=6,
             )
 
-        # ── RIGHT: Output ─────────────────────────────────────────────────────
         with gr.Column(scale=1):
             gr.Markdown("### Segmentation Output  *(zoomed to grains)*")
             seg_out = gr.Image(label="", interactive=False)
 
-    # ── Count summary ─────────────────────────────────────────────────────────
     gr.Markdown("---")
     with gr.Row():
         with gr.Column(scale=1):
@@ -439,30 +391,32 @@ with gr.Blocks(theme=THEME, title="GrainVision", css=CSS) as demo:
                 elem_id="count-box",
             )
 
-    # ── Measurement tables ────────────────────────────────────────────────────
     gr.Markdown("---")
     gr.Markdown("### Grain Measurements Table")
     with gr.Row():
         with gr.Column(scale=2):
             gr.Markdown("#### Per-Grain Measurements")
-            grain_table_out = gr.DataFrame(label="", interactive=False, wrap=False)
+            grain_table_out = gr.DataFrame(
+                label="",
+                interactive=False,
+                wrap=False,
+                height=500,          # shows all rows; scrollable if > 500px
+            )
         with gr.Column(scale=1):
             gr.Markdown("#### Summary Statistics")
-            summary_table_out = gr.DataFrame(label="", interactive=False, wrap=False)
+            summary_table_out = gr.DataFrame(
+                label="",
+                interactive=False,
+                wrap=False,
+                height=420,          # fits all 11 summary rows without scroll
+            )
 
-    # ── Two-stage wiring ──────────────────────────────────────────────────────
     OUTPUTS = [seg_out, summary_box, count_md, grain_table_out, summary_table_out]
 
-    # Stage 1: fast — fires on click, shows image + counts + loading placeholder
     run_btn.click(
-        fn      = predict_stage1,
-        inputs  = [inp_image],
-        outputs = OUTPUTS,
+        fn=predict_stage1, inputs=[inp_image], outputs=OUTPUTS,
     ).then(
-        # Stage 2: fires immediately after stage 1 completes, fills in tables
-        fn      = predict_stage2,
-        inputs  = [inp_image],
-        outputs = OUTPUTS,
+        fn=predict_stage2, inputs=[inp_image], outputs=OUTPUTS,
     )