line cropping

app.py

@@ -1,133 +1,313 @@
-import gradio as gr
+"""
+Production-Grade Engine Bolt Hole Localizer
+============================================
+Senior CV Pipeline — targets 99% detection accuracy in production.
+
+Key improvements over naive HoughCircles:
+  1. CLAHE contrast normalisation  → handles dark/uneven lighting
+  2. Multi-scale Hough sweep       → catches holes of varying apparent size
+  3. NMS deduplication             → removes duplicate circles from multi-scale
+  4. K-means (k=2) row clustering  → robust top/bottom split (not a median hack)
+  5. Intra-row outlier rejection   → drops spurious detections via IQR on y
+  6. Best-N selection per row      → always attempts 4 per row, warns if short
+  7. Least-squares horizontal fit  → straight line through row centroids (not
+                                      just first-to-last point join)
+  8. Tight crop with configurable  → returned as separate PIL image
+     padding
+"""
+
 import cv2
 import numpy as np
-
+import gradio as gr
 from PIL import Image
+from sklearn.cluster import KMeans
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# CORE DETECTION ENGINE
+# ─────────────────────────────────────────────────────────────────────────────
+
+def _clahe_preprocess(gray: np.ndarray) -> np.ndarray:
+    """CLAHE + mild Gaussian blur → maximises circle edge contrast."""
+    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
+    enhanced = clahe.apply(gray)
+    blurred  = cv2.GaussianBlur(enhanced, (5, 5), 1.2)
+    return blurred
+
 
-def detect_bolt_holes(image: np.ndarray, extend_px: int = 40):
+def _multi_scale_hough(gray_proc: np.ndarray) -> np.ndarray:
     """
-    Steps:
-    1. Detect all circular bolt holes via HoughCircles
-    2. Cluster into top-row and bottom-row (2 horizontal bands)
-    3. Draw horizontal connector lines across each row
-    4. Extend lines and add vertical connectors → tight bounding box
-    5. Return annotated image + cropped ROI
+    Run HoughCircles across a parameter grid and merge all candidates.
+    Returns array of shape (N, 3): [cx, cy, radius].
     """
-    img_rgb = np.array(image)
-    gray = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2GRAY)
-
-    # ── Step 1: Hough Circle Detection ──────────────────────────────────────
-    circles = cv2.HoughCircles(
-        gray,
-        cv2.HOUGH_GRADIENT,
-        dp=1,
-        minDist=30,
-        param1=80,
-        param2=30,
-        minRadius=10,
-        maxRadius=35,
-    )
+    h, w = gray_proc.shape
+    # Scale minRadius / maxRadius relative to image size so the detector
+    # generalises to different resolutions.
+    short_side  = min(h, w)
+    min_r = max(8,  int(short_side * 0.012))
+    max_r = min(80, int(short_side * 0.075))
+
+    param_grid = [
+        # (dp, minDist, p1, p2)
+        (1.0, 25, 100, 28),
+        (1.0, 25,  80, 22),
+        (1.2, 28,  60, 18),
+        (1.5, 20,  50, 15),
+    ]
+
+    all_circles = []
+    for dp, min_dist, p1, p2 in param_grid:
+        c = cv2.HoughCircles(
+            gray_proc,
+            cv2.HOUGH_GRADIENT,
+            dp=dp,
+            minDist=min_dist,
+            param1=p1,
+            param2=p2,
+            minRadius=min_r,
+            maxRadius=max_r,
+        )
+        if c is not None:
+            all_circles.extend(c[0].tolist())
+
+    if not all_circles:
+        return np.empty((0, 3))
 
-    if circles is None:
-        return image, image, "❌ No circles detected. Try a clearer engine image."
+    return np.array(all_circles, dtype=np.float32)
+
+
+def _nms_circles(circles: np.ndarray, iou_thresh: float = 0.35) -> np.ndarray:
+    """
+    Non-Maximum Suppression for circles.
+    Merges overlapping detections by keeping the one with the largest radius
+    (proxy for detection confidence) and discarding neighbours whose centre
+    distance < (r1 + r2) * iou_thresh.
+    """
+    if len(circles) == 0:
+        return circles
 
-    circles = np.round(circles[0]).astype(int)
+    # Sort descending by radius (largest first → keep)
+    order  = np.argsort(-circles[:, 2])
+    keep   = []
+    used   = np.zeros(len(circles), dtype=bool)
+
+    for i in order:
+        if used[i]:
+            continue
+        keep.append(i)
+        cx, cy, cr = circles[i]
+        for j in order:
+            if used[j] or j == i:
+                continue
+            ox, oy, or_ = circles[j]
+            dist = np.hypot(cx - ox, cy - oy)
+            if dist < (cr + or_) * (1.0 - iou_thresh):
+                used[j] = True
+
+    return circles[keep]
+
+
+def _cluster_rows(circles: np.ndarray, n_rows: int = 2) -> tuple:
+    """
+    K-means on y-coordinate ��� robust top/bottom row assignment.
+    Returns (top_circles, bot_circles) each as (N,3) arrays.
+    """
+    if len(circles) < n_rows:
+        return circles, np.empty((0, 3))
+
+    ys = circles[:, 1].reshape(-1, 1)
+    km = KMeans(n_clusters=n_rows, n_init=10, random_state=42).fit(ys)
+    labels = km.labels_
+
+    # Identify which label is top (smaller y) vs bottom
+    centres_y = [ys[labels == k].mean() for k in range(n_rows)]
+    top_label  = int(np.argmin(centres_y))
+    bot_label  = 1 - top_label
+
+    top = circles[labels == top_label]
+    bot = circles[labels == bot_label]
+    return top, bot
+
+
+def _reject_outliers_iqr(row: np.ndarray, axis: int = 1) -> np.ndarray:
+    """Drop points whose y-value is an outlier within their own row (IQR rule)."""
+    if len(row) < 3:
+        return row
+    vals = row[:, axis]
+    q1, q3 = np.percentile(vals, 25), np.percentile(vals, 75)
+    iqr    = q3 - q1
+    mask   = (vals >= q1 - 1.5 * iqr) & (vals <= q3 + 1.5 * iqr)
+    return row[mask]
+
+
+def _best_n(row: np.ndarray, n: int = 4) -> np.ndarray:
+    """Return up to n circles from a row, sorted left-to-right by x."""
+    row_sorted = row[np.argsort(row[:, 0])]
+    return row_sorted[:n]
+
+
+def _fit_horizontal_line(points: np.ndarray) -> float:
+    """
+    Least-squares horizontal fit: returns the mean y of the row.
+    (For bolt holes in a flat pattern this is the correct model.)
+    """
+    return float(np.mean(points[:, 1]))
 
-    # ── Step 2: Cluster into 2 horizontal rows ──────────────────────────────
-    ys = [c[1] for c in circles]
-    y_median = np.median(ys)
-    top_candidates = [(x, y) for x, y, r in circles if y < y_median]
-    bot_candidates = [(x, y) for x, y, r in circles if y >= y_median]
 
-    # Keep only 4 best per row (sorted by x, filtered by y-spread)
-    def pick_row(pts, n=4):
-        pts = sorted(pts, key=lambda p: p[0])
-        return pts[:n] if len(pts) >= n else pts
+# ─────────────────────────────────────────────────────────────────────────────
+# VISUALISATION
+# ─────────────────────────────────────────────────────────────────────────────
 
-    top_row = pick_row(top_candidates)
-    bot_row = pick_row(bot_candidates)
-    all_holes = top_row + bot_row
+PALETTE = {
+    "hole_ring" : (255, 220,   0),   # yellow
+    "hole_dot"  : (255,  60,  60),   # red
+    "h_line"    : (  0, 230,  80),   # green  — horizontal row lines
+    "bbox"      : (255,  80,  80),   # red    — bounding box
+    "label_bg"  : ( 30,  30,  30),
+    "label_fg"  : (255, 255, 255),
+}
 
-    if len(all_holes) < 4:
-        return image, image, f"⚠️ Only {len(all_holes)} holes found — need at least 4."
 
-    # ── Step 3: Draw on canvas ───────────────────────────────────────────────
+def _draw_annotations(
+    img_rgb : np.ndarray,
+    top_row : np.ndarray,
+    bot_row : np.ndarray,
+    extend_px: int,
+) -> tuple:
+    """
+    Draw:
+      • Yellow circle rings + red centre dots on each hole
+      • ONE straight green horizontal line per row (fitted, not endpoint-joined)
+      • Red bounding rectangle
+
+    Returns (annotated_img, (x1,y1,x2,y2) crop coords).
+    """
     vis = img_rgb.copy()
-    HOLE_COLOR   = (255, 220, 0)    # yellow rings
-    DOT_COLOR    = (255, 60, 60)    # red center dots
-    HLINE_COLOR  = (0, 230, 80)     # green horizontal connectors
-    BBOX_COLOR   = (255, 80, 80)    # red bounding box
-    EXT_COLOR    = (0, 180, 255)    # cyan extended lines
-
-    # Mark each bolt hole
-    hole_r = 20
-    for x, y in all_holes:
-        cv2.circle(vis, (x, y), hole_r, HOLE_COLOR, 3)
-        cv2.circle(vis, (x, y), 5, DOT_COLOR, -1)
-
-    # Horizontal green lines between holes in each row
+    h, w = vis.shape[:2]
+
+    all_holes = np.vstack([top_row, bot_row]) if len(bot_row) else top_row
+
+    # ── Bounding box ──────────────────────────────────────────────────────────
+    all_x = all_holes[:, 0]
+    all_y = all_holes[:, 1]
+    x1 = max(0, int(all_x.min()) - extend_px)
+    y1 = max(0, int(all_y.min()) - extend_px)
+    x2 = min(w, int(all_x.max()) + extend_px)
+    y2 = min(h, int(all_y.max()) + extend_px)
+
+    # ── Horizontal lines (drawn first so holes render on top) ─────────────────
     if len(top_row) >= 2:
-        cv2.line(vis, top_row[0], top_row[-1], HLINE_COLOR, 3)
+        y_top = int(_fit_horizontal_line(top_row))
+        cv2.line(vis, (x1, y_top), (x2, y_top), PALETTE["h_line"], 3, cv2.LINE_AA)
+
     if len(bot_row) >= 2:
-        cv2.line(vis, bot_row[0], bot_row[-1], HLINE_COLOR, 3)
-
-    # ── Step 4: Extended bounding box ───────────────────────────────────────
-    all_x = [p[0] for p in all_holes]
-    all_y = [p[1] for p in all_holes]
-    x_min = min(all_x) - extend_px
-    x_max = max(all_x) + extend_px
-    y_min = min(all_y) - extend_px
-    y_max = max(all_y) + extend_px
-
-    # Cyan extended horizontal lines (full width of bbox)
-    if top_row:
-        y_top = int(np.mean([p[1] for p in top_row]))
-        cv2.line(vis, (x_min, y_top), (x_max, y_top), EXT_COLOR, 2)
-    if bot_row:
-        y_bot = int(np.mean([p[1] for p in bot_row]))
-        cv2.line(vis, (x_min, y_bot), (x_max, y_bot), EXT_COLOR, 2)
-
-    # Red bounding box (vertical connectors close the rectangle)
-    cv2.rectangle(vis, (x_min, y_min), (x_max, y_max), BBOX_COLOR, 3)
-
-    # Redraw holes on top so box doesn't cover them
-    for x, y in all_holes:
-        cv2.circle(vis, (x, y), hole_r, HOLE_COLOR, 3)
-        cv2.circle(vis, (x, y), 5, DOT_COLOR, -1)
-
-    # ── Step 5: Crop ROI ─────────────────────────────────────────────────────
-    h, w = img_rgb.shape[:2]
-    cx1 = max(0, x_min)
-    cy1 = max(0, y_min)
-    cx2 = min(w, x_max)
-    cy2 = min(h, y_max)
-    cropped = img_rgb[cy1:cy2, cx1:cx2]
+        y_bot = int(_fit_horizontal_line(bot_row))
+        cv2.line(vis, (x1, y_bot), (x2, y_bot), PALETTE["h_line"], 3, cv2.LINE_AA)
+
+    # ── Bounding box ──────────────────────────────────────────────────────────
+    cv2.rectangle(vis, (x1, y1), (x2, y2), PALETTE["bbox"], 3, cv2.LINE_AA)
+
+    # ── Hole markers ─────────────────────────────────────────────────────────
+    for i, (row, label) in enumerate([(top_row, "T"), (bot_row, "B")]):
+        for j, (cx, cy, cr) in enumerate(row):
+            cx, cy, cr = int(cx), int(cy), max(int(cr), 14)
+            cv2.circle(vis, (cx, cy), cr + 4, PALETTE["hole_ring"], 3, cv2.LINE_AA)
+            cv2.circle(vis, (cx, cy), 5,      PALETTE["hole_dot"],  -1)
+            # Small label
+            tag = f"{label}{j+1}"
+            (tw, th), _ = cv2.getTextSize(tag, cv2.FONT_HERSHEY_SIMPLEX, 0.45, 1)
+            tx, ty = cx - tw // 2, cy - cr - 8
+            cv2.rectangle(vis, (tx - 2, ty - th - 2), (tx + tw + 2, ty + 2),
+                          PALETTE["label_bg"], -1)
+            cv2.putText(vis, tag, (tx, ty), cv2.FONT_HERSHEY_SIMPLEX,
+                        0.45, PALETTE["label_fg"], 1, cv2.LINE_AA)
+
+    return vis, (x1, y1, x2, y2)
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# PUBLIC API
+# ─────────────────────────────────────────────────────────────────────────────
+
+def detect_bolt_holes(image: Image.Image, extend_px: int = 40):
+    """
+    Full production pipeline.
+    Returns: (annotated PIL, cropped PIL, stats markdown)
+    """
+    img_rgb = np.array(image.convert("RGB"))
+    gray    = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2GRAY)
+
+    # 1. Preprocess
+    gray_proc = _clahe_preprocess(gray)
+
+    # 2. Multi-scale Hough
+    raw_circles = _multi_scale_hough(gray_proc)
+    if len(raw_circles) == 0:
+        return image, image, "❌ No circles detected — try a clearer engine image."
 
+    # 3. NMS deduplication
+    circles = _nms_circles(raw_circles)
+
+    # 4. Cluster into top / bottom rows
+    top_raw, bot_raw = _cluster_rows(circles)
+
+    # 5. Outlier rejection within each row
+    top_clean = _reject_outliers_iqr(top_raw)
+    bot_clean = _reject_outliers_iqr(bot_raw)
+
+    # 6. Pick best 4 per row
+    top_row = _best_n(top_clean, n=4)
+    bot_row = _best_n(bot_clean, n=4)
+
+    total = len(top_row) + len(bot_row)
+    if total < 4:
+        return image, image, f"⚠️ Only {total} bolt holes found — need ≥ 4."
+
+    # 7. Draw + crop
+    annotated, (x1, y1, x2, y2) = _draw_annotations(img_rgb, top_row, bot_row, extend_px)
+    cropped = img_rgb[y1:y2, x1:x2]
+
+    # 8. Stats
+    def fmt_row(row):
+        return ", ".join(f"({int(x)},{int(y)})" for x, y, _ in row)
+
+    status = "✅" if total == 8 else "⚠️ Partial"
     stats = (
-        f"✅ Detected **{len(all_holes)} bolt holes**\n"
-        f"• Top row ({len(top_row)} holes): {top_row}\n"
-        f"• Bottom row ({len(bot_row)} holes): {bot_row}\n"
-        f"• Bounding Box: ({x_min}, {y_min}) → ({x_max}, {y_max})\n"
-        f"• Cropped ROI size: {cx2-cx1} × {cy2-cy1} px"
+        f"{status} Detected **{total}/8 bolt holes**\n\n"
+        f"**Top row** ({len(top_row)} holes): {fmt_row(top_row)}\n\n"
+        f"**Bottom row** ({len(bot_row)} holes): {fmt_row(bot_row)}\n\n"
+        f"**Bounding Box**: ({x1}, {y1}) → ({x2}, {y2})  "
+        f"[{x2-x1} × {y2-y1} px]\n\n"
+        f"**Raw candidates before NMS**: {len(raw_circles)} → after NMS: {len(circles)}"
     )
 
-    return Image.fromarray(vis), Image.fromarray(cropped), stats
+    return Image.fromarray(annotated), Image.fromarray(cropped), stats
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# GRADIO UI
+# ─────────────────────────────────────────────────────────────────────────────
+
+CSS = """
+.gradio-container { max-width: 1100px !important; }
+#title { text-align: center; font-family: 'Courier New', monospace; }
+"""
+
+with gr.Blocks(title="Bolt Hole Localizer", theme=gr.themes.Monochrome(), css=CSS) as demo:
 
-# ── Gradio UI ────────────────────────────────────────────────────────────────
-with gr.Blocks(
-    title="Bolt Hole Localizer – Engine CV",
-    theme=gr.themes.Default(primary_hue="blue"),
-) as demo:
     gr.Markdown(
         """
-        # 🔩 Engine Bolt Hole Localization
-        ### Senior Computer Vision Pipeline
-        Upload an engine block image. The model will:
-        1. **Detect** all 8 bolt holes (4 top + 4 bottom)
-        2. **Connect** each row with horizontal lines
-        3. **Extend** lines and add vertical connectors → tight bounding box
-        4. **Crop** and return the detected region
-        """
+        # 🔩 Engine Bolt Hole Localizer — Production CV Pipeline
+        **Upload an engine block image.** The pipeline will:
+        1. **CLAHE** contrast normalisation + Gaussian smoothing
+        2. **Multi-scale HoughCircles** sweep across 4 parameter sets
+        3. **NMS deduplication** to eliminate duplicate detections
+        4. **K-means row clustering** → robust top / bottom split
+        5. **IQR outlier rejection** within each row
+        6. **Least-squares horizontal lines** through each row centroid
+        7. **Tight bounding box** + **cropped ROI** output
+        """,
+        elem_id="title"
     )
 
     with gr.Row():
@@ -135,20 +315,18 @@ with gr.Blocks(
 
     with gr.Row():
         extend_slider = gr.Slider(
-            minimum=10, maximum=100, value=40, step=5,
-            label="Bounding Box Extension (px)"
-
-
-
+            minimum=10, maximum=120, value=40, step=5,
+            label="Bounding Box Padding (px)",
+            info="Extra pixels added on each side of the outermost holes"
         )
 
-    run_btn = gr.Button("🔍 Detect Bolt Holes", variant="primary")
+    run_btn = gr.Button("🔍 Detect Bolt Holes", variant="primary", size="lg")
 
     with gr.Row():
-        out_annotated = gr.Image(label="📌 Annotated – Holes + Bounding Box")
-        out_cropped   = gr.Image(label="✂️ Cropped ROI")
+        out_annotated = gr.Image(label="📌 Annotated — Holes + Horizontal Lines + Bounding Box")
+        out_cropped   = gr.Image(label="✂️  Cropped ROI")
 
-    out_stats = gr.Markdown(label="Detection Stats")
+    out_stats = gr.Markdown(label="Detection Report")
 
     run_btn.click(
         fn=detect_bolt_holes,
@@ -156,10 +334,5 @@ with gr.Blocks(
         outputs=[out_annotated, out_cropped, out_stats],
     )
 
-    gr.Examples(
-        examples=[["perfect1.jpeg"]],
-        inputs=inp_img,
-    )
-
 if __name__ == "__main__":
     demo.launch()