detect_cave.py

#!/usr/bin/env python3
"""
detect_cave.py — Automatic cave entrance detector for IR/NIR imagery.

Usage:
    python detect_cave.py                      # batch: all jpg/png in current dir
    python detect_cave.py img1.png img2.png    # specific images

v4 — Improved pipeline (opencv + numpy only, no external models):
  - IR physics depth map (darkness × multi-scale local uniformity) as new signal
  - Texture gate: penalises textured rock/vegetation masquerading as voids
  - Vertical centroid gate: suppresses top-of-frame artefacts
  - GrabCut boundary refinement after candidate selection
  - Contour smoothing in refine_mask (wrap-around Gaussian)
  - Amber dilation ring in result visualisation
"""

import cv2
import numpy as np
import os
import sys
import glob


# ──────────────────────────────────────────────────────────────────────────────
# 1. LOAD
# ──────────────────────────────────────────────────────────────────────────────

def load_image(path: str):
    """Load image → (gray_u8, gray_f32 [0..1])."""
    img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
    if img is None:
        raise FileNotFoundError(f"Cannot read: {path}")
    if img.ndim == 3:
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    else:
        gray = img.copy()
    gray_u8 = gray.astype(np.uint8)
    gray_f32 = gray_u8.astype(np.float32) / 255.0
    return gray_u8, gray_f32


# ──────────────────────────────────────────────────────────────────────────────
# 2. PREPROCESS
# ──────────────────────────────────────────────────────────────────────────────

def preprocess_image(gray_u8, gray_f32):
    """
    Gentle preprocessing:
      - Median denoise
      - Very-large-blur background illumination estimate
      - Division normalisation (preserves cave darkness relative to local bg)
    """
    h, w = gray_u8.shape

    denoised = cv2.medianBlur(gray_u8, 5)

    # Background: huge blur (40% of min dimension)
    bg_k = max(3, int(min(h, w) * 0.40) | 1)
    background = cv2.GaussianBlur(denoised.astype(np.float32),
                                   (bg_k, bg_k), 0)
    background = np.clip(background, 10.0, 255.0)

    # Division normalisation
    corrected_f = denoised.astype(np.float32) / background
    corrected_u8 = np.clip(corrected_f * 170, 0, 255).astype(np.uint8)

    return {
        "denoised": denoised,
        "background": background,
        "corrected_u8": corrected_u8,
    }


# ──────────────────────────────────────────────────────────────────────────────
# 3. VALID REGION
# ──────────────────────────────────────────────────────────────────────────────

def compute_valid_region(gray_f32):
    """
    Soft weight map based on horizontal illumination profile.
    Uses 80th percentile per column, smoothed.
    Returns (weight_map, left_col, right_col, profile_norm).
    """
    h, w = gray_f32.shape

    col_profile = np.percentile(gray_f32, 80, axis=0).astype(np.float32)
    smooth_k = max(5, int(w * 0.08) | 1)
    profile_smooth = cv2.GaussianBlur(
        col_profile.reshape(1, -1), (smooth_k, 1), 0
    ).flatten()

    pmax = max(profile_smooth.max(), 1e-6)
    profile_norm = profile_smooth / pmax

    drop_thresh = 0.45
    actual_left_col = 0
    for c in range(w):
        if profile_norm[c] >= drop_thresh:
            actual_left_col = c
            break
    actual_right_col = w - 1
    for c in range(w - 1, -1, -1):
        if profile_norm[c] >= drop_thresh:
            actual_right_col = c
            break

    left_col  = min(actual_left_col,  int(w * 0.30))
    right_col = max(actual_right_col, int(w * 0.70))

    # Soft weight map
    weight_row = np.ones(w, dtype=np.float32)
    for c in range(w):
        if c < left_col:
            weight_row[c] = 0.3 + 0.7 * c / max(left_col, 1)
        elif c > right_col:
            weight_row[c] = 0.3 + 0.7 * (w - 1 - c) / max(w - 1 - right_col, 1)
    weight_row *= np.clip(profile_norm / drop_thresh, 0.3, 1.0)
    weight_row = np.clip(weight_row, 0.0, 1.0)

    weight_map = np.tile(weight_row, (h, 1))
    return weight_map, left_col, right_col, profile_norm, actual_left_col, actual_right_col


# ──────────────────────────────────────────────────────────────────────────────
# 4. IR PHYSICS DEPTH
# ──────────────────────────────────────────────────────────────────────────────

def compute_ir_depth(gray_f32):
    """
    Fast IR physics depth map: darkness × local_uniformity at 3 scales.

    Cave voids absorb all IR → near-black AND very uniform.
    Textured surfaces (rock, vegetation) can appear dark but non-uniform.
    Returns a depth map in [0..1]; higher = more likely to be a deep cavity.
    """
    h, w = gray_f32.shape
    darkness = 1.0 - gray_f32
    depths = []
    for base_k in [15, 31, 61]:
        ksize = max(3, min(base_k, min(h, w) // 3) | 1)
        mean_l  = cv2.GaussianBlur(gray_f32, (ksize, ksize), 0)
        mean_sq = cv2.GaussianBlur(gray_f32 * gray_f32, (ksize, ksize), 0)
        var_l   = np.clip(mean_sq - mean_l * mean_l, 0.0, None)
        std_l   = np.sqrt(var_l)
        # uniformity: 1 when perfectly uniform, drops toward 0 with high relative std
        denom      = np.clip(mean_l + 0.05, 0.05, None)
        uniformity = 1.0 - np.clip(std_l / denom, 0.0, 1.0)
        depths.append(darkness * uniformity)
    return np.mean(depths, axis=0).astype(np.float32)


# ──────────────────────────────────────────────────────────────────────────────
# 5. GENERATE CANDIDATES
# ──────────────────────────────────────────────────────────────────────────────

def _extract_components(binary, min_area):
    """Extract connected components ≥ min_area after morphological cleaning."""
    k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))
    binary = cv2.morphologyEx(binary, cv2.MORPH_OPEN,  k)
    binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, k)
    n, labels, stats, _ = cv2.connectedComponentsWithStats(binary, 8)
    result = []
    for i in range(1, n):
        if stats[i, cv2.CC_STAT_AREA] >= min_area:
            result.append(((labels == i) * 255).astype(np.uint8))
    return result


def _extract_components_heavy(binary, min_area, h, w):
    """Extract components with HEAVY morphological bridging (large closing).
    Bridges fragmented dark spots that belong to the same cave entrance."""
    close_size = max(15, int(min(h, w) * 0.04) | 1)
    close_k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE,
                                         (close_size, close_size))
    bridged = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, close_k)
    k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))
    bridged = cv2.morphologyEx(bridged, cv2.MORPH_OPEN, k)
    n, labels, stats, _ = cv2.connectedComponentsWithStats(bridged, 8)
    result = []
    for i in range(1, n):
        if stats[i, cv2.CC_STAT_AREA] >= min_area:
            result.append(((labels == i) * 255).astype(np.uint8))
    return result


def generate_candidates(proc, gray_f32, h, w, left_col=0, right_col=None):
    """
    Multi-strategy candidate generation:
      A. Multi-level thresholding with standard cleaning
      B. Multi-level thresholding with heavy bridging
      C. Iterative seed-growth from darkest pixels
      D. Otsu thresholding
      E. Adaptive threshold intersected with a dark base
      F. Valid-zone-only masking (lateral shadows masked out)
    """
    denoised     = proc["denoised"]
    corrected_u8 = proc["corrected_u8"]

    candidates = []
    min_area = int(h * w * 0.008)  # 0.8%

    # ── A. Multi-level standard thresholding ──────────────────────────────────
    for pct in [10, 15, 20, 25, 30, 35, 40]:
        thr = int(np.percentile(denoised, pct))
        _, binary = cv2.threshold(denoised, thr, 255, cv2.THRESH_BINARY_INV)
        candidates += _extract_components(binary, min_area)

    # Same on corrected
    for pct in [15, 25, 35, 45]:
        thr = int(np.percentile(corrected_u8, pct))
        _, binary = cv2.threshold(corrected_u8, thr, 255, cv2.THRESH_BINARY_INV)
        candidates += _extract_components(binary, min_area)

    # ── B. Multi-level with HEAVY bridging ────────────────────────────────────
    for pct in [10, 15, 20, 25, 30, 35]:
        thr = int(np.percentile(denoised, pct))
        _, binary = cv2.threshold(denoised, thr, 255, cv2.THRESH_BINARY_INV)
        candidates += _extract_components_heavy(binary, min_area, h, w)

    # ── C. Iterative seed-growth ──────────────────────────────────────────────
    p1 = int(np.percentile(denoised, 1))
    _, seed = cv2.threshold(denoised, max(p1, 3), 255, cv2.THRESH_BINARY_INV)
    seed_k = cv2.getStructuringElement(
        cv2.MORPH_ELLIPSE,
        (max(7, int(min(h, w) * 0.03) | 1),
         max(7, int(min(h, w) * 0.03) | 1))
    )
    seed = cv2.morphologyEx(seed, cv2.MORPH_CLOSE, seed_k)

    for pct in [5, 10, 15, 20, 25, 30, 35, 40, 50]:
        thr = int(np.percentile(denoised, pct))
        _, dark_level = cv2.threshold(denoised, thr, 255, cv2.THRESH_BINARY_INV)
        grow_k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (15, 15))
        grown = cv2.dilate(seed, grow_k, iterations=2)
        grown = cv2.bitwise_and(grown, dark_level)
        grown = cv2.morphologyEx(grown, cv2.MORPH_CLOSE, seed_k)
        seed = cv2.bitwise_or(seed, grown)
        candidates += _extract_components(grown, min_area)

    # ── D. Otsu ───────────────────────────────────────────────────────────────
    _, th_otsu = cv2.threshold(denoised, 0, 255,
                               cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    candidates += _extract_components(th_otsu, min_area)
    candidates += _extract_components_heavy(th_otsu, min_area, h, w)

    # ── E. Adaptive + dark base ───────────────────────────────────────────────
    block = max(11, int(min(h, w) * 0.15) | 1)
    th_adapt = cv2.adaptiveThreshold(
        denoised, 255,
        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY_INV,
        blockSize=block, C=10
    )
    med_val = int(np.median(denoised))
    _, dark_base = cv2.threshold(denoised, med_val, 255, cv2.THRESH_BINARY_INV)
    combined = cv2.bitwise_and(th_adapt, dark_base)
    candidates += _extract_components_heavy(combined, min_area, h, w)

    # ── F. Valid-zone-only thresholding ──────────────────────────────────────
    if right_col is None:
        right_col = w - 1
    if left_col > 10 or right_col < w - 11:
        masked_den = denoised.copy()
        masked_den[:, :left_col] = 255
        masked_den[:, right_col+1:] = 255
        for pct in [10, 15, 20, 25, 30, 35, 40]:
            thr = int(np.percentile(denoised, pct))
            _, binary = cv2.threshold(masked_den, thr, 255, cv2.THRESH_BINARY_INV)
            candidates += _extract_components(binary, min_area)
            candidates += _extract_components_heavy(binary, min_area, h, w)

    # ── Deduplicate (IoU > 0.80) ──────────────────────────────────────────────
    unique = []
    for cand in candidates:
        cand_nz = np.count_nonzero(cand)
        is_dup = False
        for ref in unique:
            inter = np.count_nonzero(cand & ref)
            union = cand_nz + np.count_nonzero(ref) - inter
            if union > 0 and inter / union > 0.80:
                is_dup = True
                break
        if not is_dup:
            unique.append(cand)

    return unique


# ──────────────────────────────────────────────────────────────────────────────
# 6. SCORE A CANDIDATE
# ──────────────────────────────────────────────────────────────────────────────

def score_candidate(mask, gray_f32, weight_map, left_col, right_col,
                    darkest5_mask, depth_map=None):
    """
    Multi-criteria scoring with MULTIPLICATIVE gates.

    Key design:
      - Contrast vs surround is the primary additive signal
      - IR physics depth rewards dark AND uniform regions (true voids)
      - Texture gate (multiplicative) penalises textured rock/vegetation
      - Vertical gate (multiplicative) suppresses top-frame artefacts
      - Area and solidity are MULTIPLICATIVE — wrong size/shape kills score
    """
    h, w = gray_f32.shape
    img_area = h * w
    mask_bool = mask.astype(bool)
    area = int(mask_bool.sum())
    if area < 10:
        return {"total": -1.0}

    # ── Geometry ──────────────────────────────────────────────────────────────
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL,
                                    cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return {"total": -1.0}
    cnt = max(contours, key=cv2.contourArea)
    cnt_area = cv2.contourArea(cnt)
    hull_area = cv2.contourArea(cv2.convexHull(cnt))
    solidity = cnt_area / hull_area if hull_area > 0 else 0.0
    x, y, bw, bh = cv2.boundingRect(cnt)
    aspect = min(bw, bh) / max(bw, bh) if max(bw, bh) > 0 else 0.0
    area_frac = area / img_area

    # ── Intensity ─────────────────────────────────────────────────────────────
    vals_inside = gray_f32[mask_bool]
    mean_inside = float(vals_inside.mean())
    std_inside  = float(vals_inside.std())
    darkness    = 1.0 - mean_inside

    # ── 1. Contrast vs wide surround (ADDITIVE, primary) ─────────────────────
    ring_width = max(40, int(min(h, w) * 0.08))
    dil_k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE,
                                       (ring_width, ring_width))
    dilated = cv2.dilate(mask, dil_k)
    ring = dilated.astype(bool) & (~mask_bool)
    if ring.sum() > 100:
        mean_outside = float(gray_f32[ring].mean())
    else:
        mean_outside = float(gray_f32.mean())
    contrast = mean_outside - mean_inside
    contrast_score = np.clip(contrast / 0.25, 0.0, 1.0)

    # ── 2. Darkness score (ADDITIVE) ──────────────────────────────────────────
    dark_score = np.clip(darkness / 0.7, 0.0, 1.0)

    # ── 3. Darkness enrichment (ADDITIVE) ─────────────────────────────────────
    darkest5_bool = darkest5_mask.astype(bool)
    total_darkest = max(darkest5_bool.sum(), 1)
    contained_frac = float((mask_bool & darkest5_bool).sum()) / total_darkest
    enrichment = contained_frac / max(area_frac, 0.001)
    enrichment_score = np.clip((enrichment - 1.0) / 8.0, 0.0, 1.0)

    # ── 4. Distance-transform depth (ADDITIVE) ───────────────────────────────
    dist_transform = cv2.distanceTransform(mask, cv2.DIST_L2, 5)
    max_dist = float(dist_transform.max())
    ref_dist = min(h, w) * 0.12
    depth_score = np.clip(max_dist / ref_dist, 0.0, 1.0)

    # ── 5. IR physics depth (ADDITIVE) ────────────────────────────────────────
    # Cave voids are dark AND spatially uniform; textured surfaces score lower
    if depth_map is not None:
        ir_depth_score = float(np.clip(depth_map[mask_bool].mean() / 0.5, 0.0, 1.0))
    else:
        ir_depth_score = dark_score * 0.5   # fallback without depth map

    # ── 6. Boundary gradient (ADDITIVE) ───────────────────────────────────────
    grad_x = cv2.Sobel(gray_f32, cv2.CV_32F, 1, 0, ksize=5)
    grad_y = cv2.Sobel(gray_f32, cv2.CV_32F, 0, 1, ksize=5)
    grad_mag = np.sqrt(grad_x**2 + grad_y**2)
    thin_k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))
    contour_ring = cv2.dilate(mask, thin_k) - cv2.erode(mask, thin_k)
    contour_bool = contour_ring.astype(bool)
    if contour_bool.sum() > 0:
        gradient_score = np.clip(float(grad_mag[contour_bool].mean()) / 0.10,
                                  0.0, 1.0)
    else:
        gradient_score = 0.0

    # ── 7. Valid region alignment (ADDITIVE) ──────────────────────────────────
    valid_score = float(weight_map[mask_bool].mean())

    # ── 8. Aspect ratio (ADDITIVE) ────────────────────────────────────────────
    aspect_score = 1.0 if aspect >= 0.15 else aspect / 0.15

    # ── 9. Position (ADDITIVE, very mild centre bias) ─────────────────────────
    cx = x + bw / 2.0
    cy = y + bh / 2.0
    dist_x = abs(cx / w - 0.5) * 2
    dist_y = abs(cy / h - 0.5) * 2
    position_score = 1.0 - 0.10 * dist_x - 0.05 * dist_y

    # ── Additive base score ───────────────────────────────────────────────────
    # Weights sum to 1.0:
    # 0.24+0.14+0.06+0.12+0.10+0.09+0.06+0.03+0.04+0.12 = 1.00
    additive = (
        0.24 * contrast_score
      + 0.14 * dark_score
      + 0.06 * enrichment_score
      + 0.12 * depth_score          # distance-transform depth
      + 0.10 * ir_depth_score       # IR physics depth (darkness × uniformity)
      + 0.09 * gradient_score
      + 0.06 * valid_score
      + 0.03 * aspect_score
      + 0.04 * position_score
      + 0.12 * 1.0                  # base
    )

    # ── MULTIPLICATIVE GATES ──────────────────────────────────────────────────

    # Area gate: 8%–28% ideal; large blobs (>28%) are usually outdoor
    # shadows, not cave entrances — penalise them more steeply than before.
    if area_frac < 0.005:
        area_mult = 0.05
    elif area_frac < 0.02:
        area_mult = 0.05 + 0.20 * (area_frac - 0.005) / 0.015
    elif area_frac < 0.04:
        area_mult = 0.25 + 0.25 * (area_frac - 0.02) / 0.02
    elif area_frac < 0.08:
        area_mult = 0.50 + 0.50 * (area_frac - 0.04) / 0.04
    elif area_frac <= 0.28:
        area_mult = 1.0
    elif area_frac <= 0.45:
        area_mult = 1.0 - 0.80 * (area_frac - 0.28) / 0.17
    else:
        area_mult = max(0.05, 0.20 - 0.15 * (area_frac - 0.45) / 0.55)

    # Solidity gate: very non-convex (donut, tentacles) → penalised
    if solidity >= 0.45:
        solidity_mult = 1.0
    elif solidity >= 0.25:
        solidity_mult = 0.4 + 0.6 * (solidity - 0.25) / 0.20
    else:
        solidity_mult = 0.4

    # Texture gate: cave voids are dark AND uniform; textured regions penalised.
    # std_inside > 0.10 starts the ramp; above 0.22 caps at 0.60.
    if std_inside <= 0.10:
        texture_mult = 1.0
    elif std_inside <= 0.22:
        texture_mult = 1.0 - 0.40 * (std_inside - 0.10) / 0.12
    else:
        texture_mult = 0.60

    # Vertical gate: entrances in the top quarter of the frame are unlikely.
    # Penalises bright sky patches and illuminated ceiling artefacts.
    if cy / h < 0.25:
        vert_gate = 0.75
    elif cy / h < 0.35:
        vert_gate = 0.75 + 0.25 * (cy / h - 0.25) / 0.10
    else:
        vert_gate = 1.0

    # Lateral penalty
    lateral_pen = 1.0
    if int(cx) < left_col or int(cx) > right_col:
        if valid_score < 0.5:
            lateral_pen = 0.4

    # Border contact penalty: true cave entrances are contained within the frame.
    # Dark patches that bleed to an image edge are likely cast shadows or vegetation
    # extending beyond the visible area, not a cave void.
    bp = 3  # border margin in pixels
    border_touches = int(
        (mask[:bp, :] > 0).any()       # top
        + (mask[-bp:, :] > 0).any()    # bottom
        + (mask[:, :bp] > 0).any()     # left
        + (mask[:, -bp:] > 0).any()    # right
    )
    if border_touches == 0:
        border_mult = 1.0
    elif border_touches == 1:
        border_mult = 0.60
    else:
        border_mult = 0.35

    total = (additive * area_mult * solidity_mult
             * texture_mult * vert_gate * lateral_pen * border_mult)

    return {
        "total":          round(float(total), 4),
        "additive":       round(float(additive), 3),
        "contrast":       round(float(contrast_score), 3),
        "dark":           round(float(dark_score), 3),
        "enrichment":     round(float(enrichment_score), 3),
        "depth":          round(float(depth_score), 3),
        "ir_depth":       round(float(ir_depth_score), 3),
        "texture":        round(float(std_inside), 3),
        "texture_mult":   round(float(texture_mult), 3),
        "vert_gate":      round(float(vert_gate), 3),
        "area_mult":      round(float(area_mult), 3),
        "area_frac":      round(float(area_frac), 4),
        "solidity":       round(float(solidity), 3),
        "sol_mult":       round(float(solidity_mult), 3),
        "gradient":       round(float(gradient_score), 3),
        "valid_score":    round(float(valid_score), 3),
        "mean_inside":    round(float(mean_inside), 3),
        "mean_outside":   round(float(mean_outside), 3),
        "border_touches": border_touches,
        "border_mult":    round(float(border_mult), 3),
    }


# ──────────────────────────────────────────────────────────────────────────────
# 7. SELECT BEST
# ──────────────────────────────────────────────────────────────────────────────

def select_best_candidate(candidates, gray_f32, weight_map,
                          left_col, right_col, depth_map=None):
    """Score all candidates, return (best_mask, best_scores, all_scores)."""
    if not candidates:
        return None, {}, []

    p5 = np.percentile(gray_f32, 5)
    darkest5_mask = (gray_f32 <= p5).astype(np.uint8) * 255

    all_scores = []
    for cand in candidates:
        sc = score_candidate(cand, gray_f32, weight_map, left_col, right_col,
                             darkest5_mask, depth_map=depth_map)
        all_scores.append(sc)

    best_idx = max(range(len(all_scores)),
                   key=lambda i: all_scores[i]["total"])
    return candidates[best_idx], all_scores[best_idx], all_scores


# ──────────────────────────────────────────────────────────────────────────────
# 8. GRABCUT REFINE
# ──────────────────────────────────────────────────────────────────────────────

def grabcut_refine(gray_u8, mask, conservative_mask=None, expand_ratio=2.5):
    """
    Refine mask boundary using GrabCut (OpenCV graph-cut, no extra deps).

    If conservative_mask is provided (the pre-expansion baseline), it is used
    as definite FG so GrabCut anchors on the known-good core and can include
    additional dark interior pixels without over-trimming.

    Without conservative_mask: eroded core is definite FG.
    With conservative_mask: conservative_mask is definite FG; extra pixels in
    mask become probable FG, letting GrabCut decide which dark interior areas
    (e.g. cave floor below a rock band) are genuinely part of the entrance.
    """
    h, w = gray_u8.shape
    area = np.count_nonzero(mask)
    if area < 200:
        return mask

    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL,
                                    cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return mask
    cnt = max(contours, key=cv2.contourArea)
    x, y, bw, bh = cv2.boundingRect(cnt)
    if bw < 5 or bh < 5:
        return mask

    # Expand bounding rectangle
    ex = int(bw * (expand_ratio - 1) / 2)
    ey = int(bh * (expand_ratio - 1) / 2)
    x1 = max(0, x - ex);  y1 = max(0, y - ey)
    x2 = min(w, x + bw + ex);  y2 = min(h, y + bh + ey)
    if x2 - x1 < 5 or y2 - y1 < 5:
        return mask

    gc_mask = np.full((h, w), cv2.GC_BGD, dtype=np.uint8)

    # Probable FG: dilated mask clipped to expanded rect
    dil_r = max(5, int(min(bw, bh) * 0.10))
    dil_k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2*dil_r+1, 2*dil_r+1))
    prob_fg = cv2.dilate(mask, dil_k)
    prob_fg[:y1, :] = 0;  prob_fg[y2:, :] = 0
    prob_fg[:, :x1] = 0;  prob_fg[:, x2:] = 0
    gc_mask[prob_fg > 0] = cv2.GC_PR_FGD

    if conservative_mask is not None and np.count_nonzero(conservative_mask) >= 10:
        # Definite FG: the pre-expansion mask (known-good entrance core)
        gc_mask[conservative_mask > 0] = cv2.GC_FGD
    else:
        # Definite FG: eroded core of input mask
        ero_r = max(3, int(min(bw, bh) * 0.08))
        ero_k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE,
                                           (2*ero_r+1, 2*ero_r+1))
        core = cv2.erode(mask, ero_k)
        gc_mask[core > 0] = cv2.GC_FGD

    # Probable BG: inside expanded rect but well away from mask
    far_r = max(7, int(min(bw, bh) * 0.20))
    far_k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2*far_r+1, 2*far_r+1))
    far_dil = cv2.dilate(mask, far_k)
    in_rect = np.zeros((h, w), np.uint8)
    in_rect[y1:y2, x1:x2] = 255
    prob_bg = cv2.bitwise_and(in_rect, cv2.bitwise_not(far_dil))
    gc_mask[prob_bg > 0] = cv2.GC_PR_BGD

    if (gc_mask == cv2.GC_FGD).sum() < 10:
        return mask

    bgd_model = np.zeros((1, 65), np.float64)
    fgd_model = np.zeros((1, 65), np.float64)
    rect = (x1, y1, x2 - x1, y2 - y1)

    try:
        vis3 = cv2.cvtColor(gray_u8, cv2.COLOR_GRAY2BGR)
        cv2.grabCut(vis3, gc_mask, rect, bgd_model, fgd_model,
                    3, cv2.GC_INIT_WITH_MASK)
        result = np.where(
            (gc_mask == cv2.GC_FGD) | (gc_mask == cv2.GC_PR_FGD),
            255, 0
        ).astype(np.uint8)

        if np.count_nonzero(result) < area * 0.25:
            return mask

        # Keep the largest component that overlaps the original mask.
        # (Guards against GrabCut fragmenting into many tiny pieces.)
        n_comp, labels, stats, _ = cv2.connectedComponentsWithStats(result, 8)
        if n_comp > 2:
            overlap_ids = np.unique(labels[mask > 0])
            overlap_ids = overlap_ids[overlap_ids != 0]
            if len(overlap_ids) > 0:
                keep_id = overlap_ids[
                    np.argmax(stats[overlap_ids, cv2.CC_STAT_AREA])
                ]
                result = ((labels == keep_id) * 255).astype(np.uint8)

        if np.count_nonzero(result) < area * 0.25:
            return mask

        return result

    except Exception:
        return mask


# ──────────────────────────────────────────────────────────────────────────────
# 9. REFINE MASK
# ──────────────────────────────────────────────────────────────────────────────

def refine_mask(mask, gray_f32):
    """Close gaps, fill holes, smooth boundary, keep largest component."""
    h, w = gray_f32.shape
    orig_area = np.count_nonzero(mask)

    cs = max(11, int(min(h, w) * 0.02) | 1)
    ck = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (cs, cs))
    refined = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, ck)

    # Fill interior holes safely using a 1-px black border.
    bordered = np.zeros((h + 2, w + 2), np.uint8)
    bordered[1:-1, 1:-1] = refined
    flood = bordered.copy()
    pad = np.zeros((h + 4, w + 4), np.uint8)
    cv2.floodFill(flood, pad, (0, 0), 255)
    holes = cv2.bitwise_not(flood)[1:-1, 1:-1]
    refined = cv2.bitwise_or(refined, holes)

    # Smooth
    sk = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (9, 9))
    refined = cv2.morphologyEx(refined, cv2.MORPH_CLOSE, sk)
    refined = cv2.morphologyEx(refined, cv2.MORPH_OPEN,  sk)

    # Largest component only
    n, labels, stats, _ = cv2.connectedComponentsWithStats(refined, 8)
    if n > 1:
        largest = 1 + np.argmax(stats[1:, cv2.CC_STAT_AREA])
        refined = ((labels == largest) * 255).astype(np.uint8)

    # ── Contour smoothing (wrap-around Gaussian) ──────────────────────────────
    cnts, _ = cv2.findContours(refined, cv2.RETR_EXTERNAL,
                                cv2.CHAIN_APPROX_NONE)
    if cnts:
        main_cnt = max(cnts, key=cv2.contourArea)
        pts = main_cnt.reshape(-1, 2).astype(np.float32)
        n_pts = len(pts)
        if n_pts > 30:
            sigma = min(15.0, max(4.0, n_pts / 120.0))
            ksize = max(3, int(6 * sigma) | 1)
            pad_n = ksize // 2
            padded = np.concatenate([pts[-pad_n:], pts, pts[:pad_n]], axis=0)
            kernel = cv2.getGaussianKernel(ksize, sigma).flatten()
            sx = np.convolve(padded[:, 0], kernel, mode='valid')[:n_pts]
            sy = np.convolve(padded[:, 1], kernel, mode='valid')[:n_pts]
            sx = np.clip(sx, 0, w - 1)
            sy = np.clip(sy, 0, h - 1)
            smooth_cnt = (np.stack([sx, sy], axis=1)
                          .astype(np.int32).reshape(-1, 1, 2))
            smooth_mask = np.zeros_like(refined)
            cv2.fillPoly(smooth_mask, [smooth_cnt], 255)
            # Safety: don't shrink more than 30%
            if np.count_nonzero(smooth_mask) >= np.count_nonzero(refined) * 0.70:
                refined = smooth_mask

    # Safety: revert if refinement bloated the mask beyond 2× original
    if np.count_nonzero(refined) > max(orig_area * 2, h * w * 0.50):
        return mask

    return refined


# ──────────────────────────────────────────────────────────────────────────────
# 10. DRAW RESULT
# ──────────────────────────────────────────────────────────────────────────────

def draw_result(gray_u8, refined_mask, scores,
                out_path, mask_path, debug_valid_path,
                weight_map, profile_norm,
                debug_cands_path, all_candidates, all_scores):
    """Save result overlay, mask, and debug images."""
    h, w = gray_u8.shape

    # ── Main result ───────────────────────────────────────────────────────────
    vis = cv2.cvtColor(gray_u8, cv2.COLOR_GRAY2BGR)

    # Amber dilation ring — buffer zone around detected entrance
    dil_r = max(5, int(min(h, w) * 0.025))
    dil_k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2*dil_r+1, 2*dil_r+1))
    dil_mask  = cv2.dilate(refined_mask, dil_k)
    ring_mask = cv2.bitwise_and(dil_mask, cv2.bitwise_not(refined_mask))
    ring_overlay = vis.copy()
    ring_overlay[ring_mask > 0] = (30, 160, 255)   # amber (BGR)
    cv2.addWeighted(ring_overlay, 0.28, vis, 0.72, 0, vis)

    # Green cave entrance overlay
    overlay = vis.copy()
    overlay[refined_mask > 0] = (100, 210, 60)
    cv2.addWeighted(overlay, 0.35, vis, 0.65, 0, vis)

    contours, _ = cv2.findContours(refined_mask, cv2.RETR_EXTERNAL,
                                    cv2.CHAIN_APPROX_SIMPLE)
    cv2.drawContours(vis, contours, -1, (0, 255, 80), 2)

    score_val = scores.get("total", 0.0)
    label = f"cave entrance  score={score_val:.2f}"
    if contours:
        cnt = max(contours, key=cv2.contourArea)
        x, y, bw, bh = cv2.boundingRect(cnt)
        tx, ty = x + 5, max(y - 12, 25)
    else:
        tx, ty = 10, 30

    fs = max(0.55, min(w, h) / 900)
    th = max(1, int(fs * 2))
    cv2.putText(vis, label, (tx+2, ty+2), cv2.FONT_HERSHEY_SIMPLEX,
                fs, (0,0,0), th+2)
    cv2.putText(vis, label, (tx, ty), cv2.FONT_HERSHEY_SIMPLEX,
                fs, (0,255,120), th)

    cv2.imwrite(out_path, vis)
    cv2.imwrite(mask_path, refined_mask)

    # ── Debug: valid region ───────────────────────────────────────────────────
    dv = cv2.cvtColor(gray_u8, cv2.COLOR_GRAY2BGR)
    for ch in range(3):
        c = dv[:,:,ch].astype(np.float32)
        if ch == 2:
            c = c * weight_map + 180 * (1.0 - weight_map)
        else:
            c = c * weight_map
        dv[:,:,ch] = np.clip(c, 0, 255).astype(np.uint8)
    for c in range(w - 1):
        y1 = h - 1 - int(profile_norm[c]     * 59)
        y2 = h - 1 - int(profile_norm[c + 1] * 59)
        cv2.line(dv, (c, y1), (c+1, y2), (0,255,255), 1)
    cv2.putText(dv, "valid region (red=penalised)", (10,25),
                cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255,255,0), 2)
    cv2.imwrite(debug_valid_path, dv)

    # ── Debug: candidates ─────────────────────────────────────────────────────
    dc = cv2.cvtColor(gray_u8, cv2.COLOR_GRAY2BGR)
    colours = [(255,80,0),(0,80,255),(200,0,200),(0,200,200),
               (200,200,0),(0,160,80),(128,128,255),(255,128,128)]
    indexed = sorted(range(len(all_candidates)),
                     key=lambda i: all_scores[i]["total"])
    for rank, i in enumerate(indexed):
        col = colours[i % len(colours)]
        cl, _ = cv2.findContours(all_candidates[i], cv2.RETR_EXTERNAL,
                                  cv2.CHAIN_APPROX_SIMPLE)
        cv2.drawContours(dc, cl, -1, col, 1)
        if rank >= len(indexed) - 5 and cl:
            c0 = max(cl, key=cv2.contourArea)
            M = cv2.moments(c0)
            if M["m00"] > 0:
                cx_m = int(M["m10"]/M["m00"])
                cy_m = int(M["m01"]/M["m00"])
                sc_v = all_scores[i]["total"]
                cv2.putText(dc, f"{sc_v:.2f}", (cx_m, cy_m),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.4, col, 1)
    cv2.drawContours(dc, contours, -1, (255,255,255), 2)
    cv2.putText(dc,
                f"{len(all_candidates)} candidates (white=best, {score_val:.2f})",
                (10,25), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (255,255,255), 2)
    cv2.imwrite(debug_cands_path, dc)


# ──────────────────────────────────────────────────────────────────────────────
# 11. PROCESS ONE IMAGE
# ──────────────────────────────────────────────────────────────────────────────

def process_image(input_path, output_dir):
    """Full pipeline for one image."""
    bn = os.path.splitext(os.path.basename(input_path))[0]
    out_r  = os.path.join(output_dir, f"{bn}_result.png")
    out_m  = os.path.join(output_dir, f"{bn}_mask.png")
    out_dv = os.path.join(output_dir, f"{bn}_debug_valid.png")
    out_dc = os.path.join(output_dir, f"{bn}_debug_candidates.png")

    gray_u8, gray_f32 = load_image(input_path)
    h, w = gray_u8.shape
    print(f"  [{bn}] loaded {w}x{h}")

    proc      = preprocess_image(gray_u8, gray_f32)
    wmap, lc, rc, pn, actual_lc, actual_rc = compute_valid_region(gray_f32)
    depth_map = compute_ir_depth(gray_f32)
    print(f"  [{bn}] valid cols {lc}–{rc} (actual {actual_lc}–{actual_rc}, of {w})")

    candidates = generate_candidates(proc, gray_f32, h, w, lc, rc)
    print(f"  [{bn}] {len(candidates)} unique candidates")

    if not candidates:
        print(f"  [{bn}] WARNING: no candidates")
        blank = np.zeros((h, w), np.uint8)
        cv2.imwrite(out_m, blank)
        cv2.imwrite(out_r, cv2.cvtColor(gray_u8, cv2.COLOR_GRAY2BGR))
        return [out_r, out_m]

    best_mask, scores, all_sc = select_best_candidate(
        candidates, gray_f32, wmap, lc, rc, depth_map=depth_map
    )
    print(f"  [{bn}] best score {scores['total']:.3f}  "
          f"area={scores['area_frac']*100:.1f}%  "
          f"add={scores['additive']:.2f}  "
          f"contrast={scores['contrast']:.2f}  "
          f"texture={scores['texture']:.2f}(×{scores['texture_mult']:.2f})  "
          f"ir_depth={scores['ir_depth']:.2f}  "
          f"depth={scores['depth']:.2f}  "
          f"area_m={scores['area_mult']:.2f}  "
          f"sol={scores['solidity']:.2f}(×{scores['sol_mult']:.2f})  "
          f"border={scores['border_touches']}(×{scores['border_mult']:.2f})  "
          f"in={scores['mean_inside']:.2f}±{scores['texture']:.2f}  "
          f"out={scores['mean_outside']:.2f}")

    # ── Solidity filter ───────────────────────────────────────────────────────
    # Non-convex candidate (e.g. entrance merged with lateral IR shadow):
    # keep only the well-illuminated weight-map portion.
    if scores.get("solidity", 1.0) < 0.65 and np.count_nonzero(best_mask) > 100:
        _is_dark_void = scores.get("mean_inside", 1.0) < 0.15
        mask_weights = wmap[best_mask > 0]
        # Dark voids use 50th pct (gentler); others use 60th pct
        w_thresh = np.percentile(mask_weights, 50 if _is_dark_void else 60)
        high_w = ((best_mask > 0) & (wmap >= w_thresh)).astype(np.uint8) * 255
        sk = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (11, 11))
        high_w = cv2.morphologyEx(high_w, cv2.MORPH_CLOSE, sk)
        high_w = cv2.morphologyEx(high_w, cv2.MORPH_OPEN, sk)
        n_hw, labels_hw, stats_hw, centroids_hw = cv2.connectedComponentsWithStats(
            high_w, 8)
        if n_hw > 1:
            valid_comps = []
            for ci in range(1, n_hw):
                cx_ci  = centroids_hw[ci, 0]
                area_ci = stats_hw[ci, cv2.CC_STAT_AREA]
                if lc <= cx_ci <= rc and area_ci >= np.count_nonzero(best_mask) * 0.10:
                    valid_comps.append((ci, area_ci))
            if valid_comps:
                best_ci = max(valid_comps, key=lambda x: x[1])[0]
                best_mask = ((labels_hw == best_ci) * 255).astype(np.uint8)
            else:
                largest = 1 + np.argmax(stats_hw[1:, cv2.CC_STAT_AREA])
                candidate_hw = ((labels_hw == largest) * 255).astype(np.uint8)
                if np.count_nonzero(candidate_hw) >= np.count_nonzero(best_mask) * 0.15:
                    best_mask = candidate_hw

    # ── Post-selection expansion ──────────────────────────────────────────────
    # Grow selected mask into connected dark pixels at a relaxed threshold.
    # Captures the full entrance when the initial candidate covers only the core.
    best_area_frac = np.count_nonzero(best_mask) / (h * w)
    if best_area_frac < 0.25:
        relax_pct = min(50, max(30, int(scores.get("area_frac", 0.1) * 100 * 4)))
        relax_thr = int(np.percentile(proc["denoised"], relax_pct))
        _, relax_dark = cv2.threshold(proc["denoised"], relax_thr, 255,
                                       cv2.THRESH_BINARY_INV)
        br_k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE,
                                          (max(9, int(min(h,w)*0.02)|1),
                                           max(9, int(min(h,w)*0.02)|1)))
        relax_dark = cv2.morphologyEx(relax_dark, cv2.MORPH_CLOSE, br_k)
        n_rd, labels_rd, _, _ = cv2.connectedComponentsWithStats(relax_dark, 8)
        overlap_labels = set(np.unique(labels_rd[best_mask > 0])) - {0}
        if overlap_labels:
            expanded = np.zeros_like(best_mask)
            for lb in overlap_labels:
                expanded[labels_rd == lb] = 255
            # Clip to valid columns to prevent re-introducing lateral dark zones
            if lc > int(w * 0.05):
                expanded[:, :lc] = 0
            if rc < int(w * 0.95):
                expanded[:, rc+1:] = 0
            n_exp, labels_exp, stats_exp, _ = cv2.connectedComponentsWithStats(
                expanded, 8)
            if n_exp > 1:
                largest_exp = 1 + np.argmax(stats_exp[1:, cv2.CC_STAT_AREA])
                expanded = ((labels_exp == largest_exp) * 255).astype(np.uint8)
            exp_area_frac = np.count_nonzero(expanded) / (h * w)
            if exp_area_frac <= 0.40 and exp_area_frac > best_area_frac * 0.8:
                exp_mean  = float(gray_f32[expanded > 0].mean())
                orig_mean = float(gray_f32[best_mask > 0].mean())
                # Reject expansion if centroid drifted far — guards against
                # bridging to a disconnected dark zone (e.g. vegetation corner).
                orig_pts = np.argwhere(best_mask > 0).astype(np.float32)
                exp_pts  = np.argwhere(expanded > 0).astype(np.float32)
                orig_cy_m, orig_cx_m = orig_pts.mean(axis=0)
                exp_cy_m,  exp_cx_m  = exp_pts.mean(axis=0)
                centroid_shift = (np.sqrt((exp_cx_m - orig_cx_m)**2
                                         + (exp_cy_m - orig_cy_m)**2)
                                  / min(h, w))
                if exp_mean < orig_mean + 0.15 and centroid_shift <= 0.20:
                    print(f"  [{bn}] expanded {best_area_frac*100:.1f}% → "
                          f"{exp_area_frac*100:.1f}%")
                    best_mask = expanded

    # ── GrabCut boundary refinement ───────────────────────────────────────────
    pre_gc = np.count_nonzero(best_mask) / (h * w)
    gc_result = grabcut_refine(gray_u8, best_mask, expand_ratio=2.0)
    post_gc = np.count_nonzero(gc_result) / (h * w)
    if post_gc > 0:
        print(f"  [{bn}] grabcut {pre_gc*100:.1f}% → {post_gc*100:.1f}%")
        best_mask = gc_result

    refined = refine_mask(best_mask, gray_f32)

    # Hard-clip final mask to valid illumination columns.
    # Removes penumbra/vignetting zones from the result even when the initial
    # candidate or GrabCut extended into them.
    if lc > int(w * 0.05):
        refined[:, :lc] = 0
    if rc < int(w * 0.95):
        refined[:, rc + 1:] = 0

    draw_result(gray_u8, refined, scores,
                out_r, out_m, out_dv,
                wmap, pn,
                out_dc, candidates, all_sc)

    final_area = np.count_nonzero(refined) / (h * w)
    print(f"  [{bn}] final area {final_area*100:.1f}%")
    outputs = [out_r, out_m, out_dv, out_dc]
    for p in outputs:
        print(f"  [{bn}] saved: {os.path.basename(p)}")
    return outputs


# ──────────────────────────────────────────────────────────────────────────────
# 12. MAIN
# ──────────────────────────────────────────────────────────────────────────────

def main():
    if len(sys.argv) >= 2:
        # One or more explicit image paths
        for img_path in sys.argv[1:]:
            out_dir = os.path.dirname(os.path.abspath(img_path)) or "."
            print(f"Processing: {os.path.basename(img_path)}")
            process_image(img_path, out_dir)
            print()
    else:
        # Batch mode: process all jpg/png in the script's directory
        cwd = os.path.dirname(os.path.abspath(__file__))
        patterns = ["*.jpg","*.jpeg","*.png","*.JPG","*.JPEG","*.PNG"]
        found = []
        for pat in patterns:
            found.extend(glob.glob(os.path.join(cwd, pat)))
        suffixes = ("_result.png","_mask.png","_debug_valid.png",
                    "_debug_candidates.png")
        inputs = sorted(set(
            f for f in found
            if not any(os.path.basename(f).endswith(s) for s in suffixes)
        ))
        if not inputs:
            print("No input images found.")
            sys.exit(1)
        print(f"Found {len(inputs)} input image(s):")
        for p in inputs:
            print(f"  {os.path.basename(p)}")
        print()
        all_out = []
        for img in inputs:
            print(f"Processing: {os.path.basename(img)}")
            all_out += process_image(img, cwd)
            print()
        print("─" * 60)
        print(f"Done. {len(all_out)} output files:")
        for p in all_out:
            print(f"  {os.path.basename(p)}")


if __name__ == "__main__":
    main()