sp.py

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()