pipeline.py

"""
pipeline.py — Conveyor Belt Damage Detection (Classical CV)
"""

import os
import json
import argparse
import cv2
import numpy as np
from scipy.ndimage import label as scipy_label, uniform_filter


# ─────────────────────────────────────────────
#  CONFIGURATION (tune these if needed)
# ─────────────────────────────────────────────
CFG = {
    # Blur kernel for residual computation (must be odd)
    "blur_ksize": 61,

    # Scratch: percentile threshold when belt is highly damaged (belt_median > 8)
    # Lowered 88→85 to catch more scratches in night/high-texture images
    "scratch_pct_high": 85,
    # Scratch: percentile threshold when belt is clean/slightly damaged
    # Lowered 95→92 to catch more scratches in day images
    "scratch_pct_low": 92,

    # Minimum blob area (px) to be counted as a scratch region
    # Raised 5000→8000 to reduce small false-positive boxes
    "scratch_min_area": 8000,
    # Minimum density (fraction of bbox pixels that are anomalous)
    # Raised 0.05→0.07 to reduce false positives
    "scratch_min_density": 0.07,

    # Edge strip width as fraction of belt width
    # Raised 0.08→0.12: edge damage often spans wider than 8%
    "edge_strip_frac": 0.12,
    # Minimum edge strip width in pixels
    "edge_strip_min_px": 50,
    # Edge damage: ratio of edge texture to center texture to flag
    # Lowered 2.5→1.4: night images have high baseline texture so ratio was never reaching 2.5
    "edge_ratio_thresh": 1.4,
    # Minimum damaged-row length (px) to include in edge bbox
    "edge_min_row_run": 30,

    # Mask erosion size to avoid polygon-boundary artifacts
    "mask_erosion": 20,

    # Night brightness threshold (mean belt pixel < this → night)
    "night_brightness_thresh": 100,

    # CLAHE clip limits
    "clahe_night": 3.0,
    "clahe_day": 2.0,
}


# ─────────────────────────────────────────────
#  LABEL PARSING
# ─────────────────────────────────────────────
def parse_polygon(label_line: str, img_w: int, img_h: int) -> np.ndarray:
    """
    Parse a YOLO-format polygon label line.
    Format: <class_id> x1 y1 x2 y2 ... xn yn  (all normalized 0-1)
    Returns: numpy array of (N, 2) integer pixel coordinates.
    """
    vals = label_line.strip().split()
    coords = [float(v) for v in vals[1:]]
    points = [
        (int(coords[i] * img_w), int(coords[i + 1] * img_h))
        for i in range(0, len(coords) - 1, 2)
    ]
    return np.array(points, dtype=np.int32)


def load_label(label_path: str, img_w: int, img_h: int):
    """
    Load label file and return polygon points for the first annotation.
    Returns None if file is missing or empty.
    """
    if not os.path.exists(label_path):
        return None
    with open(label_path, "r") as f:
        lines = [l.strip() for l in f if l.strip()]
    if not lines:
        return None
    # Use the first (and typically only) annotation
    return parse_polygon(lines[0], img_w, img_h)


def fallback_belt_mask(img_w: int, img_h: int) -> tuple:
    """
    If no label is available, assume belt occupies the middle 70% of image width.
    """
    x_left = int(img_w * 0.15)
    x_right = int(img_w * 0.85)
    pts = np.array([
        [x_left, 0],
        [x_right, 0],
        [x_right, img_h],
        [x_left, img_h],
    ], dtype=np.int32)
    return pts


# ─────────────────────────────────────────────
#  BELT MASK
# ─────────────────────────────────────────────
def build_masks(pts: np.ndarray, img_w: int, img_h: int):
    """
    Build full belt mask and eroded inner mask (avoids boundary artifacts).
    Returns: full_mask, inner_mask (both uint8 0/255)
    """
    full_mask = np.zeros((img_h, img_w), dtype=np.uint8)
    cv2.fillPoly(full_mask, [pts], 255)

    erosion_size = CFG["mask_erosion"]
    kernel = cv2.getStructuringElement(
        cv2.MORPH_RECT, (erosion_size, erosion_size)
    )
    inner_mask = cv2.erode(full_mask, kernel)
    return full_mask, inner_mask


# ─────────────────────────────────────────────
#  PREPROCESSING
# ─────────────────────────────────────────────
def preprocess(img: np.ndarray, inner_mask: np.ndarray):
    """
    Convert to grayscale, apply CLAHE for adaptive contrast enhancement.
    Also returns: is_night flag, belt_median residual value.
    """
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    belt_pixels = gray[inner_mask > 0]
    belt_brightness = belt_pixels.mean() if len(belt_pixels) > 0 else 128.0

    is_night = belt_brightness < CFG["night_brightness_thresh"]
    clip = CFG["clahe_night"] if is_night else CFG["clahe_day"]
    clahe = cv2.createCLAHE(clipLimit=clip, tileGridSize=(8, 8))
    enhanced = clahe.apply(gray)

    # Compute residual (local anomaly map)
    ksize = CFG["blur_ksize"]
    img_f = enhanced.astype(np.float32)
    blur = cv2.GaussianBlur(img_f, (ksize, ksize), 0)
    residual = np.abs(img_f - blur)

    belt_res_vals = residual[inner_mask > 0]
    belt_median = float(np.median(belt_res_vals)) if len(belt_res_vals) > 0 else 0.0

    return enhanced, residual, is_night, belt_median


# ─────────────────────────────────────────────
#  SCRATCH DETECTION
# ─────────────────────────────────────────────
def detect_scratches(
    residual: np.ndarray,
    inner_mask: np.ndarray,
    belt_median: float,
) -> list:
    """
    Detect scratch damage on belt surface using residual anomaly map.
    Scratches appear as elongated high-residual regions.

    FIX 1: Tiered adaptive threshold based on belt_median.
    Manual review showed 83% FP rate on clean images using a single threshold.
    Root cause: clean belts have low median residual but local texture
    variations still cross a fixed percentile threshold.
    Solution: use belt_median to classify belt cleanliness and apply
    progressively stricter thresholds for cleaner belts.

    Returns list of [x_min, y_min, x_max, y_max] bboxes.
    """
    belt_res_vals = residual[inner_mask > 0]
    if len(belt_res_vals) == 0:
        return []

    # ── Tiered threshold (FIX 1) ─────────────────────────────────────
    # belt_median < 3.0 → belt is very clean → use pct=97 (near-zero detections)
    # belt_median 3.0-6.0 → borderline → use pct=93
    # belt_median > 6.0 → clearly damaged (scratch/texture present) → use pct=85
    if belt_median < 3.0:
        pct = 97
    elif belt_median < 6.0:
        pct = 93
    else:
        pct = CFG["scratch_pct_high"]   # 85

    scratch_thresh = float(np.percentile(belt_res_vals, pct))

    scratch_binary = (residual > scratch_thresh) & (inner_mask > 0)
    scratch_u8 = scratch_binary.astype(np.uint8) * 255

    # Morphological closing: connect nearby scratch fragments
    kh = cv2.getStructuringElement(cv2.MORPH_RECT, (101, 7))
    kv = cv2.getStructuringElement(cv2.MORPH_RECT, (7, 101))
    closed = cv2.morphologyEx(scratch_u8, cv2.MORPH_CLOSE, kh)
    closed = cv2.morphologyEx(closed, cv2.MORPH_CLOSE, kv)

    # Morphological opening: remove tiny noise blobs
    k_open = cv2.getStructuringElement(cv2.MORPH_RECT, (21, 21))
    scratch_clean = cv2.morphologyEx(closed, cv2.MORPH_OPEN, k_open)

    # Extract bboxes from connected components
    labeled_map, n_components = scipy_label(scratch_clean > 0)
    bboxes = []
    min_area    = CFG["scratch_min_area"]
    min_density = CFG["scratch_min_density"]

    for comp_id in range(1, n_components + 1):
        region = labeled_map == comp_id
        if region.sum() < min_area:
            continue

        ys, xs = np.where(region)
        x1, y1 = int(xs.min()), int(ys.min())
        x2, y2 = int(xs.max()), int(ys.max())

        # Width filter: very narrow boxes (< 50px) are almost always FP
        if (x2 - x1) < 50:
            continue

        # ── FIX 1A: Aspect ratio filter ──────────────────────────────
        # Manual review (136 images) showed:
        #   True scratch bbox aspect (h/w) median = 3.6  → tall, elongated
        #   False positive aspect (h/w) median   = 2.1  → short, squarish
        # Reject boxes that are not elongated enough to be real scratches.
        bbox_h  = y2 - y1
        bbox_w  = x2 - x1
        aspect  = bbox_h / (bbox_w + 1e-6)
        if aspect < 2.0:
            continue

        # ── FIX 1B: Minimum height filter ────────────────────────────
        # Real scratches span a significant portion of belt height.
        # Short boxes (< 200px on a 2160px image) are almost always noise.
        if bbox_h < 200:
            continue

        # Density check: raw anomaly pixels / bbox area
        bbox_area = (x2 - x1 + 1) * (y2 - y1 + 1)
        raw_count = int(scratch_binary[y1 : y2 + 1, x1 : x2 + 1].sum())
        if bbox_area > 0 and raw_count / bbox_area < min_density:
            continue

        bboxes.append([x1, y1, x2, y2])

    return bboxes


# ─────────────────────────────────────────────
#  EDGE DAMAGE DETECTION
# ─────────────────────────────────────────────
def detect_edge_damage(
    residual: np.ndarray,
    inner_mask: np.ndarray,
    pts: np.ndarray,
    img_w: int,
    img_h: int,
    enhanced: np.ndarray,
) -> list:
    """
    Detect edge damage using a dual-signal approach.

    FIX 2: Previous approach (residual ratio only) had recall of ~5%.
    Root cause: night images have uniformly high residual everywhere,
    so edge/centre ratio stays low even when edge is damaged.

    New dual-signal approach:
    Signal A — Residual ratio: edge residual vs centre residual (works for day)
    Signal B — Gradient spike ratio: p90/p50 of row-wise Sobel-X gradient in
               the edge strip (works for both day and night).
               Edge damage = certain rows have much higher gradient than others
               (jagged/torn edge). Clean edge = uniform gradient across rows.

    A detection is triggered if EITHER signal exceeds its threshold.
    This improves recall while Signal A's precision gate limits FP rate.

    Returns list of [x_min, y_min, x_max, y_max] bboxes.
    """
    xs_pts = [int(p[0]) for p in pts]
    belt_left  = min(xs_pts)
    belt_right = max(xs_pts)
    belt_w_px  = belt_right - belt_left

    ew = max(int(belt_w_px * CFG["edge_strip_frac"]), CFG["edge_strip_min_px"])

    # Central reference: residual median of middle 24% of belt
    cx = (belt_left + belt_right) // 2
    half_center = int(belt_w_px * 0.12)
    center_strip = np.zeros((img_h, img_w), dtype=np.uint8)
    center_strip[:, max(0, cx - half_center): min(img_w, cx + half_center)] = 255
    center_mask   = cv2.bitwise_and(inner_mask, center_strip)
    center_res    = residual[center_mask > 0]
    center_median = float(np.median(center_res)) if len(center_res) > 0 else 1.0

    # ── FIX 2: Rewritten edge detection ──────────────────────────
    # Previous approach (residual ratio) had 5% recall on 136 images.
    # Root cause: edge damage = LOCAL irregularity at belt boundary.
    # The ratio approach compared WHOLE edge strip to centre — too coarse.
    #
    # New approach — 3 complementary signals per row in edge strip:
    #   S1: Row residual z-score vs strip's own baseline (per-strip adaptive)
    #   S2: Row Sobel-X gradient vs strip's own baseline
    #   S3: Row local std dev vs strip's own baseline
    # A row is "damaged" if ANY 2 of 3 signals fire.
    # This is much more sensitive to local irregularities.

    sobelx   = cv2.Sobel(enhanced, cv2.CV_64F, 1, 0, ksize=3)
    grad_abs = np.abs(sobelx)

    # Local std dev map
    img_f        = enhanced.astype(np.float32)
    lmean        = uniform_filter(img_f, size=15)
    lsq          = uniform_filter(img_f**2, size=15)
    local_std    = np.sqrt(np.maximum(lsq - lmean**2, 0))

    edge_bboxes = []

    for side, x_s, x_e in [
        ("L", belt_left + 10,       belt_left + ew + 10),
        ("R", belt_right - ew - 10, belt_right - 10),
    ]:
        x_s = max(0, int(x_s))
        x_e = min(img_w, int(x_e))
        if x_s >= x_e:
            continue

        strip = np.zeros((img_h, img_w), dtype=np.uint8)
        strip[:, x_s:x_e] = 255
        emask = cv2.bitwise_and(inner_mask, strip)
        if emask.sum() < 50:
            continue

        # Compute per-row signal values
        row_res = np.zeros(img_h, dtype=np.float32)   # residual
        row_grd = np.zeros(img_h, dtype=np.float32)   # gradient
        row_std = np.zeros(img_h, dtype=np.float32)   # local std

        for r in range(img_h):
            pm = emask[r, x_s:x_e]
            if pm.sum() == 0:
                continue
            row_res[r] = float(residual[r,    x_s:x_e][pm > 0].mean())
            row_grd[r] = float(grad_abs[r,    x_s:x_e][pm > 0].mean())
            row_std[r] = float(local_std[r,   x_s:x_e][pm > 0].mean())

        valid = (row_res > 0)
        if valid.sum() < 10:
            continue

        # Per-strip baseline (median of valid rows)
        base_res = float(np.median(row_res[valid]))
        base_grd = float(np.median(row_grd[valid]))
        base_std = float(np.median(row_std[valid]))

        # Adaptive thresholds: 1.6x baseline (lower = more sensitive)
        thr_res = base_res * 1.6
        thr_grd = base_grd * 1.6
        thr_std = base_std * 1.6

        # Row is damaged if at least 2 of 3 signals fire
        s1 = (row_res > thr_res) & valid
        s2 = (row_grd > thr_grd) & valid
        s3 = (row_std > thr_std) & valid
        damage_rows = ((s1.astype(int) + s2.astype(int) + s3.astype(int)) >= 2)

        # Only proceed if there's a meaningful amount of damage rows
        if damage_rows.sum() < CFG["edge_min_row_run"]:
            continue

        # Collect contiguous damaged runs, merging gaps < 300px
        # (edge damage is often scattered across rows with small gaps)
        merged_runs = []
        in_run      = False
        run_start   = 0
        for r in range(img_h):
            if damage_rows[r] and not in_run:
                in_run    = True
                run_start = r
            elif not damage_rows[r] and in_run:
                in_run = False
                length = r - run_start
                if length >= 10:
                    if merged_runs and (run_start - merged_runs[-1][1]) < 300:
                        merged_runs[-1] = (merged_runs[-1][0], r)
                    else:
                        merged_runs.append((run_start, r))
        if in_run:
            length = img_h - run_start
            if length >= 10:
                if merged_runs and (run_start - merged_runs[-1][1]) < 300:
                    merged_runs[-1] = (merged_runs[-1][0], img_h)
                else:
                    merged_runs.append((run_start, img_h))

        # Only keep runs with minimum final length
        for r_start, r_end in merged_runs:
            if (r_end - r_start) >= CFG["edge_min_row_run"]:
                edge_bboxes.append([x_s, int(r_start), x_e, r_end])

    return edge_bboxes


# ─────────────────────────────────────────────
#  NMS (remove redundant bboxes)
# ─────────────────────────────────────────────
def iou(a, b):
    ax1, ay1, ax2, ay2 = a
    bx1, by1, bx2, by2 = b
    ix1 = max(ax1, bx1); iy1 = max(ay1, by1)
    ix2 = min(ax2, bx2); iy2 = min(ay2, by2)
    if ix2 <= ix1 or iy2 <= iy1:
        return 0.0
    inter = (ix2 - ix1) * (iy2 - iy1)
    area_a = (ax2 - ax1) * (ay2 - ay1)
    area_b = (bx2 - bx1) * (by2 - by1)
    return inter / (area_a + area_b - inter + 1e-6)


def nms(bboxes: list, iou_thresh: float = 0.5) -> list:
    if not bboxes:
        return []
    # Sort by area descending
    bboxes = sorted(bboxes, key=lambda b: (b[2]-b[0])*(b[3]-b[1]), reverse=True)
    kept = []
    suppressed = [False] * len(bboxes)
    for i in range(len(bboxes)):
        if suppressed[i]:
            continue
        kept.append(bboxes[i])
        for j in range(i + 1, len(bboxes)):
            if not suppressed[j] and iou(bboxes[i], bboxes[j]) > iou_thresh:
                suppressed[j] = True
    return kept


# ─────────────────────────────────────────────
#  VISUALISATION
# ─────────────────────────────────────────────
def draw_detections(img: np.ndarray, scratch_bboxes: list, edge_bboxes: list) -> np.ndarray:
    """
    Draw bounding boxes on image copy.
    scratch → red,  edge_damage → green
    """
    vis = img.copy()
    font = cv2.FONT_HERSHEY_SIMPLEX

    for bb in scratch_bboxes:
        x1, y1, x2, y2 = bb
        cv2.rectangle(vis, (x1, y1), (x2, y2), (0, 0, 255), 3)
        label_y = max(y1 - 10, 30)
        cv2.putText(vis, "scratch", (x1, label_y), font, 1.2, (0, 0, 255), 3)

    for bb in edge_bboxes:
        x1, y1, x2, y2 = bb
        cv2.rectangle(vis, (x1, y1), (x2, y2), (0, 255, 0), 3)
        label_y = max(y1 - 10, 60)
        cv2.putText(vis, "edge_damage", (x1, label_y), font, 1.2, (0, 255, 0), 3)

    return vis


# ─────────────────────────────────────────────
#  MAIN PROCESSING FUNCTION
# ─────────────────────────────────────────────
def process_image(img_path: str, label_dir: str, output_dir: str) -> dict:
    """
    Process a single image: detect damage, save annotated image + JSON.
    Returns detection dict.
    """
    img = cv2.imread(img_path)
    if img is None:
        print(f"  [WARN] Could not read: {img_path}")
        return {}

    img_h, img_w = img.shape[:2]
    basename = os.path.basename(img_path)
    stem = os.path.splitext(basename)[0]

    # ── Load belt polygon from label file ──
    label_path = os.path.join(label_dir, stem + ".txt")
    pts = load_label(label_path, img_w, img_h)
    if pts is None:
        print(f"  [INFO] No label for {basename}, using fallback belt region")
        pts = fallback_belt_mask(img_w, img_h)

    # ── Build masks ──
    full_mask, inner_mask = build_masks(pts, img_w, img_h)

    # ── Preprocess ──
    enhanced, residual, is_night, belt_median = preprocess(img, inner_mask)

    # ── Detect ──
    scratch_bboxes = detect_scratches(residual, inner_mask, belt_median)
    edge_bboxes = detect_edge_damage(residual, inner_mask, pts, img_w, img_h, enhanced)

    # ── NMS to remove duplicate/overlapping boxes ──
    scratch_bboxes = nms(scratch_bboxes, iou_thresh=0.4)
    edge_bboxes = nms(edge_bboxes, iou_thresh=0.5)

    # ── Build detection dict ──
    detections = {}
    det_idx = 1
    for bb in scratch_bboxes:
        detections[str(det_idx)] = {
            "bbox_coordinates": [bb[0], bb[1], bb[2], bb[3]],
            "class": "scratch",
        }
        det_idx += 1
    for bb in edge_bboxes:
        detections[str(det_idx)] = {
            "bbox_coordinates": [bb[0], bb[1], bb[2], bb[3]],
            "class": "edge_damage",
        }
        det_idx += 1

    # ── Save annotated image ──
    vis = draw_detections(img, scratch_bboxes, edge_bboxes)
    out_img_path = os.path.join(output_dir, basename)
    cv2.imwrite(out_img_path, vis)

    # ── Save JSON ──
    json_data = {
        k: {"bbox_coordinates": v["bbox_coordinates"]}
        for k, v in detections.items()
    }
    out_json_path = os.path.join(output_dir, stem + ".json")
    with open(out_json_path, "w") as f:
        json.dump(json_data, f, indent=2)

    mode = "NIGHT" if is_night else "DAY"
    print(
        f"  [{mode}] {basename}: "
        f"scratch={len(scratch_bboxes)}, edge={len(edge_bboxes)}"
    )
    return detections


# ─────────────────────────────────────────────
#  CLI ENTRY POINT
# ─────────────────────────────────────────────
def main():
    parser = argparse.ArgumentParser(
        description="Conveyor Belt Damage Detection Pipeline"
    )
    parser.add_argument(
        "--image_dir",
        required=True,
        help="Path to folder containing input images (.jpg / .png)",
    )
    parser.add_argument(
        "--output_dir",
        required=True,
        help="Path to folder where annotated images and JSON will be saved",
    )
    parser.add_argument(
        "--label_dir",
        default=None,
        help="Path to YOLO label folder (optional; defaults to image_dir/../labels)",
    )
    args = parser.parse_args()

    image_dir = args.image_dir
    output_dir = args.output_dir

    # Auto-detect label directory
    if args.label_dir:
        label_dir = args.label_dir
    else:
        label_dir = os.path.join(os.path.dirname(image_dir.rstrip("/\\")), "labels")

    os.makedirs(output_dir, exist_ok=True)

    # Collect images
    valid_exts = {".jpg", ".jpeg", ".png", ".bmp"}
    image_files = sorted(
        f for f in os.listdir(image_dir)
        if os.path.splitext(f)[1].lower() in valid_exts
    )

    if not image_files:
        print(f"[ERROR] No images found in: {image_dir}")
        return

    print(f"Found {len(image_files)} images.")
    print(f"Labels dir : {label_dir}")
    print(f"Output dir : {output_dir}")
    print("-" * 60)

    total_scratch = 0
    total_edge = 0

    for fname in image_files:
        img_path = os.path.join(image_dir, fname)
        dets = process_image(img_path, label_dir, output_dir)
        s = sum(1 for v in dets.values() if v.get("class") == "scratch")
        e = sum(1 for v in dets.values() if v.get("class") == "edge_damage")
        total_scratch += s
        total_edge += e

    print("-" * 60)
    print(f"Done. Total: {total_scratch} scratch detections, {total_edge} edge_damage detections.")
    print(f"Results saved to: {output_dir}")


if __name__ == "__main__":
    main()