import gradio as gr import cv2 import numpy as np from PIL import Image from scipy.spatial import KDTree from itertools import combinations # ────────────────────────────────────────────────────────────────────────────── # CORE DETECTION HELPERS # ────────────────────────────────────────────────────────────────────────────── def preprocess(gray: np.ndarray): """Multi-scale CLAHE + bilateral denoise for robust circle detection.""" clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8)) eq = clahe.apply(gray) denoised = cv2.bilateralFilter(eq, 9, 75, 75) return denoised def detect_circles_multi(gray: np.ndarray): """ Try a sweep of HoughCircles parameters; return the best set of circles. 'Best' = most circles detected within a sane radius range. Also tries on inverted image (dark holes on bright background). """ h, w = gray.shape min_r = max(8, int(min(h, w) * 0.012)) max_r = max(40, int(min(h, w) * 0.08)) param_grid = [ # dp, minDist, p1, p2 (1, 30, 80, 28), (1, 25, 70, 22), (1, 35, 90, 32), (1.2, 30, 80, 25), (1, 20, 60, 18), (1, 40, 100, 35), ] best = None for invert in (False, True): src = (255 - gray) if invert else gray preprocessed = preprocess(src) blurred = cv2.GaussianBlur(preprocessed, (5, 5), 1.5) for dp, minDist, p1, p2 in param_grid: circles = cv2.HoughCircles( blurred, cv2.HOUGH_GRADIENT, dp=dp, minDist=minDist, param1=p1, param2=p2, minRadius=min_r, maxRadius=max_r, ) if circles is not None: c = np.round(circles[0]).astype(int) if best is None or len(c) > len(best): best = c return best # shape (N, 3): x, y, r — or None def remove_duplicate_circles(circles, tol=20): """Merge circles whose centres are within `tol` pixels.""" if len(circles) == 0: return circles kept = [] used = np.zeros(len(circles), bool) # sort by radius descending so we keep the most prominent order = np.argsort(-circles[:, 2]) for i in order: if used[i]: continue dists = np.hypot(circles[:, 0] - circles[i, 0], circles[:, 1] - circles[i, 1]) mask = dists < tol kept.append(circles[i]) used[mask] = True return np.array(kept) # ────────────────────────────────────────────────────────────────────────────── # CONNECTIVITY (works for ANY layout – horizontal, vertical, diagonal, grid) # ────────────────────────────────────────────────────────────────────────────── def build_adjacency(centres, k_neighbours=2, angle_snap_deg=10): """ Connect each hole to its k nearest neighbours, then prune to keep only 'axis-aligned' edges (horizontal / vertical / 45°). Returns list of (i, j) index pairs. """ n = len(centres) if n < 2: return [] pts = np.array(centres, dtype=float) k = min(k_neighbours + 1, n) tree = KDTree(pts) _, idxs = tree.query(pts, k=k) # each row: [self, nbr1, nbr2, …] edges = set() for i, row in enumerate(idxs): for j in row[1:]: # skip self (index 0) a, b = min(i, j), max(i, j) edges.add((a, b)) # ── angle snapping: keep edges that are close to 0°/45°/90°/135° ──────── snapped = [] snap_angles = np.array([0, 45, 90, 135, 180]) for i, j in edges: dx = pts[j, 0] - pts[i, 0] dy = pts[j, 1] - pts[i, 1] ang = np.degrees(np.arctan2(abs(dy), abs(dx))) % 180 # 0–180 diff = np.min(np.abs(snap_angles - ang)) if diff <= angle_snap_deg: snapped.append((i, j)) # Fallback: if snapping kills everything, keep raw nearest-neighbour edges return snapped if snapped else list(edges) def min_bounding_rect(pts_xy): """Axis-aligned bounding rectangle for a set of (x,y) points.""" pts = np.array(pts_xy) x0, y0 = pts[:, 0].min(), pts[:, 1].min() x1, y1 = pts[:, 0].max(), pts[:, 1].max() return int(x0), int(y0), int(x1), int(y1) # ────────────────────────────────────────────────────────────────────────────── # MAIN PIPELINE # ────────────────────────────────────────────────────────────────────────────── def detect_bolt_holes(image: Image.Image, extend_px: int = 40, k_neighbours: int = 2, angle_snap: int = 15): img_rgb = np.array(image) gray = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2GRAY) h, w = gray.shape # ── 1. Detect circles ──────────────────────────────────────────────────── raw = detect_circles_multi(gray) if raw is None: return image, image, "❌ No circles detected. Try a clearer image." circles = remove_duplicate_circles(raw, tol=20) # Filter to reasonable bolt-hole sizes (skip tiny noise and huge regions) min_r = max(8, int(min(h, w) * 0.012)) max_r = max(40, int(min(h, w) * 0.08)) circles = np.array([c for c in circles if min_r <= c[2] <= max_r]) if len(circles) < 2: return image, image, f"⚠️ Only {len(circles)} valid circle(s) found – need ≥ 2." centres = [(int(c[0]), int(c[1])) for c in circles] radii = [int(c[2]) for c in circles] # ── 2. Build adjacency graph ────────────────────────────────────────────── edges = build_adjacency(centres, k_neighbours=k_neighbours, angle_snap_deg=angle_snap) # ── 3. Draw ─────────────────────────────────────────────────────────────── vis = img_rgb.copy() RING_COLOR = (255, 220, 0) # yellow DOT_COLOR = ( 30, 200, 255) # cyan centre dot LINE_COLOR = ( 0, 230, 80) # green connector line BBOX_COLOR = (255, 60, 60) # red bounding box # Draw connector lines first (under circles) for i, j in edges: cv2.line(vis, centres[i], centres[j], LINE_COLOR, 3, cv2.LINE_AA) # Draw circles on top for (cx, cy), r in zip(centres, radii): draw_r = max(r, 14) cv2.circle(vis, (cx, cy), draw_r, RING_COLOR, 3, cv2.LINE_AA) cv2.circle(vis, (cx, cy), 5, DOT_COLOR, -1, cv2.LINE_AA) # ── 4. Bounding box ─────────────────────────────────────────────────────── x0, y0, x1, y1 = min_bounding_rect(centres) x0 -= extend_px; y0 -= extend_px x1 += extend_px; y1 += extend_px # clamp to image bounds x0 = max(0, x0); y0 = max(0, y0) x1 = min(w, x1); y1 = min(h, y1) cv2.rectangle(vis, (x0, y0), (x1, y1), BBOX_COLOR, 3) # Redraw circles so bbox doesn't overlap them for (cx, cy), r in zip(centres, radii): draw_r = max(r, 14) cv2.circle(vis, (cx, cy), draw_r, RING_COLOR, 3, cv2.LINE_AA) cv2.circle(vis, (cx, cy), 5, DOT_COLOR, -1, cv2.LINE_AA) # ── 5. Crop ROI ─────────────────────────────────────────────────────────── cropped = img_rgb[y0:y1, x0:x1] # ── 6. Stats ────────────────────────────────────────────────────────────── stats = ( f"✅ **{len(circles)} bolt holes** detected | **{len(edges)} connections** drawn\n\n" f"| # | Centre (x, y) | Radius |\n" f"|---|--------------|--------|\n" + "\n".join(f"| {i+1} | {cx}, {cy} | {r} px |" for i, ((cx, cy), r) in enumerate(zip(centres, radii))) + f"\n\n• **Bounding box**: ({x0}, {y0}) → ({x1}, {y1})\n" f"• **Crop size**: {x1-x0} × {y1-y0} px" ) return Image.fromarray(vis), Image.fromarray(cropped), stats # ────────────────────────────────────────────────────────────────────────────── # 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 ### Adaptive Computer Vision Pipeline Works with **any hole layout** – horizontal rows, vertical columns, diagonal, mixed or irregular arrangements. **Pipeline:** 1. Multi-parameter Hough + CLAHE pre-processing 2. Duplicate removal & radius filtering 3. k-NN adjacency graph → angle-snapped connector lines 4. Axis-aligned bounding box + cropped ROI """ ) with gr.Row(): inp_img = gr.Image(type="pil", label="📷 Input Engine Image") with gr.Row(): extend_slider = gr.Slider(10, 120, value=40, step=5, label="Bounding Box Extension (px)") k_slider = gr.Slider(1, 4, value=2, step=1, label="Connections per Hole (k-neighbours)") angle_slider = gr.Slider(5, 45, value=15, step=5, label="Angle Snap Tolerance (°)") run_btn = gr.Button("🔍 Detect & Connect Bolt Holes", variant="primary") with gr.Row(): out_annotated = gr.Image(label="📌 Annotated – Holes + Connections + BBox") out_cropped = gr.Image(label="✂️ Cropped ROI") out_stats = gr.Markdown(label="Detection Stats") run_btn.click( fn=detect_bolt_holes, inputs=[inp_img, extend_slider, k_slider, angle_slider], outputs=[out_annotated, out_cropped, out_stats], ) if __name__ == "__main__": demo.launch()