predict.py

#!/usr/bin/env python3
"""Find wall boundaries, fixtures, attached obstacles, and wallpaper; estimate coverage."""
import argparse
import json
from pathlib import Path

import cv2
import numpy as np
import torch
from scipy import ndimage

from completion import completion_summary
from condition_cues import condition_summary, defect_fraction_curve, rust_cue
from depth_fusion import measure_with_depth
from paint_color import MATCH_DELTA_E, color_summary
from site_objects import DEFAULT_MIN_CONFIDENCE as DEFAULT_EQUIPMENT_CONFIDENCE
from site_objects import OBJECT_RGB, detect_objects, object_summary
from wall_measure import measure_wall, paint_estimate
from label_schema import CLASSES, MATERIAL_CLASSES
from models import build_from_checkpoint, checkpoint_material_classes
from regions import damage_regions, rescale
from wallpaper import wallpaper_cues
from window_postprocess import regularized_mask, window_instances

NAMES = CLASSES
COLORS_BGR = np.array([
    [90, 40, 20],    # other/non-wall
    [230, 80, 20],   # unpainted wall
    [30, 200, 30],   # painted wall
    [0, 210, 245],   # uncertain wall
    [220, 40, 190],  # skirting/trim
    [0, 0, 255],     # switch/outlet
    [255, 255, 0],   # AC unit
    [0, 128, 255],   # door
    [255, 0, 128],   # window
    [128, 128, 0],   # generic wall-attached obstacle
], dtype=np.uint8)


def boundary_samples(wall, skirting, wall_probability=None):
    h, w = wall.shape
    ys, xs = np.nonzero(wall)
    if len(xs) == 0:
        return {"bbox_xyxy": None, "skirting_bbox_xyxy": None, "bottom_edge_samples": [],
                "wall_columns_detected_fraction": 0.0, "median_wall_to_skirting_gap_px": None}
    xmin, xmax = int(xs.min()), int(xs.max())
    samples = []
    gaps = []
    for x in np.linspace(xmin, xmax, min(12, xmax - xmin + 1), dtype=int):
        wall_col = np.flatnonzero(wall[:, x])
        skirt_col = np.flatnonzero(skirting[:, x])
        if len(wall_col):
            item = {"x": int(x), "top_y": int(wall_col[0]), "bottom_y": int(wall_col[-1])}
            if wall_probability is not None:
                item["bottom_edge_wall_probability"] = round(float(wall_probability[wall_col[-1], x]), 4)
            if len(skirt_col):
                item["skirting_top_y"] = int(skirt_col[0])
                gap = int(skirt_col[0] - wall_col[-1] - 1)
                item["wall_to_skirting_gap_px"] = gap
                gaps.append(gap)
            samples.append(item)
    sy, sx = np.nonzero(skirting)
    return {
        "bbox_xyxy": [xmin, int(ys.min()), xmax, int(ys.max())],
        "skirting_bbox_xyxy": [int(sx.min()), int(sy.min()), int(sx.max()), int(sy.max())] if len(sx) else None,
        "bottom_edge_samples": samples,
        "wall_columns_detected_fraction": float(np.count_nonzero(wall.any(axis=0)) / w),
        "median_wall_to_skirting_gap_px": float(np.median(gaps)) if gaps else None,
    }


def object_components(mask, class_id, confidence, min_pixels=8, min_confidence=0.5):
    count, _, stats, centers = cv2.connectedComponentsWithStats((mask == class_id).astype(np.uint8), 8)
    objects = []
    for i in range(1, count):
        x, y, width, height, area = map(int, stats[i])
        if area >= min_pixels:
            component = mask[y:y + height, x:x + width] == class_id
            score = float(confidence[y:y + height, x:x + width][component].mean())
            if score < min_confidence:
                continue
            objects.append({"bbox_xyxy": [x, y, x + width - 1, y + height - 1],
                            "center_xy": [round(float(centers[i, 0]), 1), round(float(centers[i, 1]), 1)],
                            "pixels": area, "mean_pixel_confidence": round(score, 4)})
    return objects


def wallpaper_summary(wall, wallpaper, probability, cues, args):
    """Combine the material head and the model-free pattern/seam cue into one wallpaper verdict.

    ``wallpaper_detected`` means the model saw it; ``possible_wallpaper`` means only
    the cue did (or the checkpoint predates the wallpaper class) - worth a close-up
    before painting, not proof. Board-and-batten siding is a known cue false positive.
    """
    total_wall = int(wall.sum())
    model_fraction = int(wallpaper.sum()) / total_wall if total_wall and probability is not None else None
    cue = cues.get("score")
    if not total_wall:
        status = "no_wall_detected"
    elif model_fraction is not None and model_fraction >= args.min_wallpaper_fraction:
        status = "wallpaper_detected"
    elif cue is not None and cue >= args.wallpaper_cue_threshold:
        status = "possible_wallpaper"
    else:
        status = "none_detected"
    return {
        "status": status,
        "model_supports_wallpaper": probability is not None,
        "fraction_of_wall": model_fraction,
        "mean_probability_on_wall": float(probability[wall].mean()) if probability is not None and total_wall else None,
        "cue": cues,
        "min_fraction": args.min_wallpaper_fraction,
        "cue_threshold": args.wallpaper_cue_threshold,
    }


def ndimage_label_count(mask):
    return ndimage.label(mask, structure=np.ones((3, 3), bool))[1]


# Uncalibrated fallback: any confident damage on 0.5% of the wall. On real photos this
# flags nearly every frame; calibrate_condition.py fits an operating point per checkpoint.
DEFAULT_CONDITION = (0.5, 0.005)


def condition_thresholds(args, checkpoint):
    """``(min_confidence, min_fraction, source)``: explicit flags, else the checkpoint's calibration, else defaults."""
    point = checkpoint.get("condition_operating_point") or {}
    confidence = getattr(args, "min_condition_confidence", None)
    fraction = getattr(args, "min_defect_fraction", None)
    source = "flags" if confidence is not None or fraction is not None else (
        "checkpoint_calibration" if point else "uncalibrated_default")
    if confidence is None:
        confidence = point.get("min_confidence", DEFAULT_CONDITION[0])
    if fraction is None:
        fraction = point.get("min_fraction", DEFAULT_CONDITION[1])
    return float(confidence), float(fraction), source


def analyze(bgr, model, checkpoint, device, args):
    """Run the model and post-processing on one BGR frame; returns the summary dict and the masks.

    ``args`` carries the thresholds ``main`` parses (any namespace with the same
    attributes works), so evaluation scripts can score many images with one loaded model.
    """
    material_names = checkpoint_material_classes(checkpoint)
    # Checkpoints trained before the wallpaper class still run; wallpaper then rests on the cue alone.
    wallpaper_index = material_names.index("wallpaper") if "wallpaper" in material_names else None
    rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
    size = int(checkpoint["size"])
    h, w = rgb.shape[:2]
    scale = size / max(h, w)
    out_w, out_h = max(1, round(w * scale)), max(1, round(h * scale))
    resized = cv2.resize(rgb, (out_w, out_h), interpolation=cv2.INTER_AREA if scale < 1 else cv2.INTER_LINEAR)
    left, top = (size - out_w) // 2, (size - out_h) // 2
    canvas = np.full((size, size, 3), 114, dtype=np.uint8)
    canvas[top:top + out_h, left:left + out_w] = resized
    x = canvas.astype(np.float32) / 255.0
    mean, std = np.asarray(checkpoint["mean"], np.float32), np.asarray(checkpoint["std"], np.float32)
    x = torch.from_numpy(((x - mean) / std).transpose(2, 0, 1).copy()).unsqueeze(0).to(device)
    with torch.inference_mode():
        output = model(x)
        if args.tta:
            flipped_output = model(torch.flip(x, dims=[3]))
            output = {name: 0.5 * (value + flipped_output[name].flip(dims=[3])) for name, value in output.items()}
        # Temperatures fitted on validation (train.py) make confidence thresholds meaningful.
        small_prob = torch.softmax(output["semantic"] / float(checkpoint.get("temperature", 1.0)), dim=1)[0].cpu().numpy()
        small_drywall_prob = torch.softmax(output["drywall"] / float(checkpoint.get("material_temperature", 1.0)),
                                           dim=1)[0].cpu().numpy()
        small_condition_prob = (torch.softmax(output["condition"], dim=1)[0].cpu().numpy()
                                if "condition" in output else None)
    small_mask = small_prob.argmax(axis=0).astype(np.uint8)
    cropped_mask = small_mask[top:top + out_h, left:left + out_w]
    mask = cv2.resize(cropped_mask, (w, h), interpolation=cv2.INTER_NEAREST)
    small_wall_prob = small_prob[1:4].sum(axis=0)
    small_wall_conf = small_prob[1:4].max(axis=0)
    small_max_conf = small_prob.max(axis=0)
    painted_probability = cv2.resize(small_prob[2][top:top + out_h, left:left + out_w], (w, h),
                                     interpolation=cv2.INTER_LINEAR)
    wall_probability = cv2.resize(small_wall_prob[top:top + out_h, left:left + out_w], (w, h), interpolation=cv2.INTER_LINEAR)
    wall_class_confidence = cv2.resize(small_wall_conf[top:top + out_h, left:left + out_w], (w, h), interpolation=cv2.INTER_LINEAR)
    max_confidence = cv2.resize(small_max_conf[top:top + out_h, left:left + out_w], (w, h), interpolation=cv2.INTER_LINEAR)
    small_material_class = small_drywall_prob.argmax(axis=0).astype(np.uint8)
    small_material_conf = small_drywall_prob.max(axis=0)
    material_class = cv2.resize(small_material_class[top:top + out_h, left:left + out_w], (w, h), interpolation=cv2.INTER_NEAREST)
    drywall_confidence = cv2.resize(small_drywall_prob[1, top:top + out_h, left:left + out_w], (w, h), interpolation=cv2.INTER_LINEAR)
    material_confidence = cv2.resize(small_material_conf[top:top + out_h, left:left + out_w], (w, h), interpolation=cv2.INTER_LINEAR)
    wallpaper_probability = None
    if wallpaper_index is not None:
        wallpaper_probability = cv2.resize(small_drywall_prob[wallpaper_index, top:top + out_h, left:left + out_w],
                                           (w, h), interpolation=cv2.INTER_LINEAR)
    low_confidence_wall = (wall_probability >= args.min_wall_probability) & (wall_class_confidence < args.min_wall_class_confidence)
    mask[low_confidence_wall] = 3

    # Regularise windows: merge mullion fragments and fit oriented rectangles so
    # keep-out geometry is complete (raw components can leave gaps over mullions).
    window_probability = cv2.resize(small_prob[8][top:top + out_h, left:left + out_w], (w, h), interpolation=cv2.INTER_LINEAR)
    window_mask = window_probability >= args.min_object_confidence
    windows, _ = window_instances(window_mask, confidence=window_probability, min_area=args.window_min_area,
                                  merge_gap=args.window_merge_gap, min_fill=args.window_min_fill,
                                  max_aspect=args.window_max_aspect)
    windows_regularized = regularized_mask(windows, (h, w)) if args.regularize_windows else window_mask
    window_missed = windows_regularized & ~np.isin(mask, [8])
    mask[window_missed] = 8

    wall = np.isin(mask, [1, 2, 3])
    painted = mask == 2
    unpainted = mask == 1
    uncertain = mask == 3
    skirting = mask == 4
    drywall_known = (material_confidence >= args.min_drywall_confidence) & wall
    drywall = (material_class == 1) & drywall_known
    drywall_unknown = (~drywall_known) & wall
    wallpaper = ((material_class == wallpaper_index) & drywall_known) if wallpaper_index is not None else np.zeros_like(wall)
    known_wall = painted | unpainted
    condition_class = condition_confidence = None
    if small_condition_prob is not None:
        condition_class = cv2.resize(small_condition_prob.argmax(axis=0).astype(np.uint8)[top:top + out_h, left:left + out_w],
                                     (w, h), interpolation=cv2.INTER_NEAREST)
        condition_confidence = cv2.resize(small_condition_prob.max(axis=0)[top:top + out_h, left:left + out_w],
                                          (w, h), interpolation=cv2.INTER_LINEAR)
    min_confidence, min_fraction, condition_source = condition_thresholds(args, checkpoint)
    condition, defect = condition_summary(known_wall, condition_class, condition_confidence, rust_cue(rgb, known_wall),
                                          min_confidence=min_confidence, min_fraction=min_fraction,
                                          cue_threshold=args.rust_cue_threshold)
    condition["thresholds_from"] = condition_source
    if small_condition_prob is not None and condition["status"] == "defects_detected":
        # Whole patches (hysteresis + merge) at the model's resolution, so pixel sizes do not
        # depend on the photo's; the grown patch joins the keep-out (it still needs scraping).
        small_known = cv2.resize(known_wall.astype(np.uint8), (out_w, out_h), interpolation=cv2.INTER_NEAREST) > 0
        found, region_mask = damage_regions(small_condition_prob[1:3, top:top + out_h, left:left + out_w],
                                            small_known, high=min_confidence,
                                            prep_margin_px=args.prep_margin_px * out_w / w)
        condition["regions"] = rescale(found, w / out_w, h / out_h)
        condition["regions_total"] = int(ndimage_label_count(region_mask))
        defect = defect | (cv2.resize(region_mask.astype(np.uint8), (w, h), interpolation=cv2.INTER_NEAREST) > 0) & known_wall
    else:
        condition["regions"], condition["regions_total"] = [], 0
    if condition["status"] in ("none_detected", "possible_rust") and condition_class is not None \
            and (min_confidence, min_fraction) != DEFAULT_CONDITION:
        # Below the calibrated point but enough for the old any-damage rule: say so, do not drop it.
        weak = defect_fraction_curve(known_wall, condition_class, condition_confidence, (DEFAULT_CONDITION[0],))
        if weak[DEFAULT_CONDITION[0]] >= DEFAULT_CONDITION[1]:
            condition["weak_damage_signal"] = round(weak[DEFAULT_CONDITION[0]], 5)
    # Peeling paint and rust need preparing before any coat: keep the brush off them.
    hazards = (mask == 0) | uncertain | (mask >= 4) | defect
    if args.keepout_margin_px:
        diameter = 2 * args.keepout_margin_px + 1
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (diameter, diameter))
        hazards = cv2.dilate(hazards.astype(np.uint8), kernel) > 0
    safe_paintable = known_wall & ~hazards
    total_wall = int(wall.sum())
    decided = int((painted | unpainted).sum())
    n_painted, n_unpainted, n_uncertain, n_skirting = map(int, (painted.sum(), unpainted.sum(), uncertain.sum(), skirting.sum()))
    summary = {
        "painted_fraction_of_decided_wall": n_painted / decided if decided else None,
        "painted_fraction_lower_bound": n_painted / total_wall if total_wall else None,
        "painted_fraction_upper_bound": (n_painted + n_uncertain) / total_wall if total_wall else None,
        "wall_uncertain_fraction": n_uncertain / total_wall if total_wall else None,
        "mean_wall_probability": float(wall_probability[wall].mean()) if total_wall else None,
        "mean_wall_class_confidence": float(wall_class_confidence[wall].mean()) if total_wall else None,
        "low_confidence_wall_pixels": int(low_confidence_wall.sum()),
        "inspection_status": "no_wall_detected" if not total_wall else ("reobserve_low_confidence" if n_uncertain / total_wall > args.max_uncertain_fraction else "coverage_estimated"),
        "wall_pixels": total_wall,
        "safe_paintable_pixels": int(safe_paintable.sum()),
        "keepout_margin_px": args.keepout_margin_px,
        "skirting_pixels": n_skirting,
        "wall_fixtures": {
            name: object_components(mask, class_id, max_confidence, args.min_object_pixels, args.min_object_confidence)
            for class_id, name in ((5, "switch_outlet"), (6, "ac_unit"), (7, "door"), (9, "wall_obstacle"))
        },
        "windows": windows,
        "windows_count": len(windows),
        "window_pixels_regularized": int(windows_regularized.sum()),
        "wall_geometry": boundary_samples(wall, skirting, wall_probability),
        "classes": {name: int(np.count_nonzero(mask == i)) for i, name in enumerate(NAMES)},
        "drywall": {
            "confident_drywall_fraction_of_known_wall": int(drywall.sum()) / max(int(wall.sum() - drywall_unknown.sum()), 1),
            "fraction_lower_bound": int(drywall.sum()) / total_wall if total_wall else None,
            "fraction_upper_bound": (int(drywall.sum()) + int(drywall_unknown.sum())) / total_wall if total_wall else None,
            "unknown_material_fraction": int(drywall_unknown.sum()) / total_wall if total_wall else None,
            "mean_drywall_probability_on_wall": float(drywall_confidence[wall].mean()) if total_wall else None,
            "mean_material_confidence_on_wall": float(material_confidence[wall].mean()) if total_wall else None,
            "confidence_threshold": args.min_drywall_confidence,
        },
        "wallpaper": wallpaper_summary(wall, wallpaper, wallpaper_probability, wallpaper_cues(rgb, wall), args),
        "surface_condition": condition,
    }
    if getattr(args, "measure", True):
        summary["measurements"] = measurements(mask, rgb, summary, args)
    completion, confident_painted, touch_up = completion_summary(
        wall, painted, painted_probability, checkpoint.get("completion_threshold"), uncertain=uncertain,
        min_region_px=args.min_touchup_pixels)
    summary["completion"] = completion
    masks = {"mask": mask, "wall": wall, "skirting": skirting, "drywall": drywall, "wallpaper": wallpaper,
             "safe_paintable": safe_paintable, "hazards": hazards, "uncertain": uncertain,
             "drywall_known": drywall_known, "material_class": material_class,
             "confident_painted": confident_painted, "touch_up": touch_up, "completion": completion,
             "condition_class": condition_class, "condition_confidence": condition_confidence,
             # P(paint failure), P(rust) at model resolution (letterbox removed), for region evaluation.
             "condition_probability_small": (small_condition_prob[1:3, top:top + out_h, left:left + out_w]
                                             if small_condition_prob is not None else None),
             "known_wall": known_wall, "defect": defect, "equipment": np.zeros_like(wall)}
    reference = getattr(args, "reference_color", None)

    # What colour is on the wall, and does it match what was specified?
    if getattr(args, "colors", True):
        summary["paint_color"] = color_summary(
            rgb, painted, faces=(summary.get("measurements") or {}).get("faces"), reference=reference,
            k=getattr(args, "color_palette_k", 3), tolerance=getattr(args, "color_tolerance", MATCH_DELTA_E))

    # What equipment is on site, is any of it standing where paint is about to go, and does the
    # tin on site carry the specified colour? Detections inside the paint area join the keep-out
    # so the planner routes around a customer's tin instead of spraying it.
    if getattr(args, "objects", True):
        tools = [tool for tool in str(getattr(args, "required_tools", "") or "").replace(" ", "").split(",") if tool]
        detections = detect_objects(rgb, min_confidence=getattr(args, "min_equipment_confidence",
                                                                DEFAULT_EQUIPMENT_CONFIDENCE))
        summary["site_objects"] = object_summary(
            detections, rgb=rgb, paint_zone=safe_paintable, required_tools=tools, reference=reference,
            min_confidence=getattr(args, "min_equipment_confidence", DEFAULT_EQUIPMENT_CONFIDENCE))
        if getattr(args, "objects_in_keepout", True) and summary["site_objects"]["in_paint_zone"]:
            equipment = np.zeros_like(wall)
            for item in summary["site_objects"]["in_paint_zone"]:
                x0, y0, x1, y1 = item["bbox_xyxy"]
                equipment[max(0, y0 - 2):y1 + 3, max(0, x0 - 2):x1 + 3] = True
            masks["equipment"] = equipment
            masks["hazards"] = hazards | equipment
            masks["safe_paintable"] = safe_paintable & ~equipment
            summary["safe_paintable_pixels"] = int(masks["safe_paintable"].sum())
            summary["equipment_keepout_pixels"] = int(equipment.sum())
            summary["site_objects"]["notes"].append(
                f"{int(equipment.sum())} px of equipment keep-out added to the mask the planner uses "
                "(--no-objects-in-keepout to allow painting over it)")
    return summary, masks


def draw_measurements(overlay, measured):
    """Stop lines on the overlay: ceiling (cyan), paint bottom (magenta), floor (grey), ends (yellow; red = no stop)."""
    thick = max(2, overlay.shape[1] // 400)
    for face in measured.get("faces", []):
        lines = face["stop_lines_px"]
        for key, colour in (("top", (255, 255, 0)), ("bottom", (255, 0, 255)), ("floor", (160, 160, 160))):
            (x0, y0), (x1, y1) = lines[key]
            cv2.line(overlay, (int(x0), int(y0)), (int(x1), int(y1)), colour, thick)
        for side in ("left", "right"):
            (x0, y0), (x1, y1) = lines[side]
            colour = (0, 0, 255) if face["ends"][side] == "frame" else (0, 255, 255)
            cv2.line(overlay, (int(x0), int(y0)), (int(x1), int(y1)), colour, thick)
        cx = int(np.mean([p[0] for p in face["corners_px"]]))
        cy = int(np.mean([p[1] for p in face["corners_px"]]))
        text = f"{face['width_ft']:.1f} x {face['paint_height_ft']:.1f} ft  {face['area_ft2']['wall_visible']:.0f} sq ft"
        scale = max(0.5, overlay.shape[1] / 1600)
        cv2.putText(overlay, text, (max(cx - 160, 5), cy), cv2.FONT_HERSHEY_SIMPLEX, scale, (0, 0, 0), thick + 2)
        cv2.putText(overlay, text, (max(cx - 160, 5), cy), cv2.FONT_HERSHEY_SIMPLEX, scale, (255, 255, 255), thick)


def draw_site_objects(overlay, summary):
    """Box the equipment and chip the sampled colour, so the report can be checked at a glance."""
    height, width = overlay.shape[:2]
    thick = max(2, width // 400)
    scale = max(0.4, width / 1600)
    for item in (summary.get("site_objects") or {}).get("detections") or []:
        x0, y0, x1, y1 = item["bbox_xyxy"]
        colour = tuple(int(v) for v in reversed(OBJECT_RGB.get(item["class"], (255, 0, 0))))
        cv2.rectangle(overlay, (x0, y0), (x1, y1), colour, thick)
        label = f"{item['class'].replace('_', ' ')} {item['confidence']:.2f}"
        cv2.putText(overlay, label, (x0 + 2, max(12, y0 - 4)), cv2.FONT_HERSHEY_SIMPLEX, scale, (0, 0, 0), thick + 2)
        cv2.putText(overlay, label, (x0 + 2, max(12, y0 - 4)), cv2.FONT_HERSHEY_SIMPLEX, scale, colour, thick)
    # The sampled coat (and the specified colour under it) as a chip on the frame.
    paint_colour = (summary.get("paint_color") or {}).get("dominant")
    if paint_colour:
        size = max(48, width // 14)
        chips = [(paint_colour, f"{paint_colour['name']} {paint_colour['hex']}")]
        match = (summary.get("paint_color") or {}).get("reference")
        if match:
            chips.append((match["reference"], f"spec {match['verdict']} (dE {match['delta_e']})"))
        for index, (entry, text) in enumerate(chips):
            y = 12 + index * (size // 2 + 6)
            cv2.rectangle(overlay, (12, y), (12 + size, y + size // 2),
                          tuple(int(v) for v in reversed(entry["rgb"])), -1)
            cv2.rectangle(overlay, (12, y), (12 + size, y + size // 2), (30, 30, 30), 1)
            cv2.putText(overlay, text, (18 + size, y + size // 3), cv2.FONT_HERSHEY_SIMPLEX, scale, (0, 0, 0), thick + 2)
            cv2.putText(overlay, text, (18 + size, y + size // 3), cv2.FONT_HERSHEY_SIMPLEX, scale, (255, 255, 255), thick)


def exif_focal_px(path, width, height):
    """Focal length in pixels from EXIF FocalLengthIn35mmFilm (36 mm = the long side), or None."""
    try:
        from PIL import Image
        with Image.open(path) as image:
            exif = image.getexif()
            value = exif.get(0xA405) or exif.get_ifd(0x8769).get(0xA405)
        return float(value) / 36.0 * max(width, height) if value else None
    except Exception:
        return None


def measurements(mask, rgb, summary, args, gray=None, principal=None, frame_size=None):
    """Stop lines, corners, wall area (sq ft) and paint needed, from the predicted mask.

    ``gray``/``principal``/``frame_size`` are given for a multi-frame mosaic, whose
    canvas extends past the reference camera frame (see predict_multi.py).
    """
    gray = rgb.mean(axis=2) / 255.0 if gray is None else gray
    measured, depth_note = None, None
    depth = getattr(args, "depth_map", None)
    if depth is not None:
        h, w = mask.shape
        intrinsics = args.intrinsics or ((args.focal_px, args.focal_px, w / 2, h / 2) if args.focal_px else None)
        if intrinsics is None:
            depth_note = "depth given without --intrinsics or EXIF focal length; measured from the image instead"
        else:
            measured = measure_with_depth(mask, depth, intrinsics)
            if measured.get("status") != "measured":
                depth_note = f"depth fusion failed ({measured.get('reason')}); measured from the image instead"
                measured = None
    if measured is None:
        measured = measure_wall(mask, gray=gray, wall_height_ft=args.wall_height_ft, wall_width_ft=args.wall_width_ft,
                                mm_per_px=args.mm_per_px, focal_px=args.focal_px, ceiling_ft=args.ceiling_height_ft,
                                principal=principal, frame_size=frame_size)
        if depth_note:
            measured["depth_note"] = depth_note
    faces = measured["faces"]
    if not faces:
        return measured
    area = {k: round(sum(f["area_ft2"][k] for f in faces), 1) for k in faces[0]["area_ft2"]}
    remaining = area["unpainted"] + area["uncertain"]
    wallpaper = (summary.get("wallpaper") or {}).get("status") == "wallpaper_detected"
    common = {"coats": args.coats, "surface": args.surface, "waste": args.waste, "application": args.application,
              "wallpaper": wallpaper}
    measured["total_area_ft2"] = area
    measured["paint"] = {
        # Every visible wall surface gets the full coats; bare (unpainted) area is primed first.
        "whole_wall": paint_estimate(area["wall_visible"], unpainted_ft2=area["unpainted"],
                                     hidden_ft2=area["hidden_behind_objects"], **common),
        # Only what is not yet confidently painted: the robot's remaining job.
        "remaining": paint_estimate(remaining, unpainted_ft2=area["unpainted"], **common),
    }
    return measured


def build_parser(positional=True, add_help=True):
    """predict.py's options; ``positional=False`` gives them to another tool as a parent parser."""
    ap = argparse.ArgumentParser(description=__doc__, add_help=add_help)
    if positional:
        ap.add_argument("image", type=Path, help="full camera frame; model locates the wall within it")
    ap.add_argument("--checkpoint", default="artifacts/best.pt",
                    help="local .pt, or hf://owner/repo[@revision] (fetches config.json, then the weights)")
    ap.add_argument("--overlay", type=Path, default=Path("painting_prediction.png"))
    ap.add_argument("--mask", type=Path, default=Path("painting_mask.png"))
    ap.add_argument("--paintable-mask", type=Path, default=Path("paintable_wall_mask.png"))
    ap.add_argument("--keepout-mask", type=Path, default=Path("fixture_keepout_mask.png"))
    ap.add_argument("--uncertainty-mask", type=Path, default=Path("wall_uncertainty_mask.png"))
    ap.add_argument("--drywall-mask", type=Path, default=Path("drywall_material_mask.png"))
    ap.add_argument("--painted-mask", type=Path, default=Path("confident_painted_mask.png"),
                    help="binary mask of confidently painted wall; feed to robot_planner.py --observe")
    ap.add_argument("--touchup-mask", type=Path, default=Path("touchup_mask.png"),
                    help="binary mask of wall not yet confidently painted (largest regions)")
    ap.add_argument("--min-touchup-pixels", type=int, default=64,
                    help="smallest not-done region reported for touch-up")
    ap.add_argument("--json", type=Path, help="optional path to save the summary JSON")
    ap.add_argument("--min-wall-probability", type=float, default=0.45)
    ap.add_argument("--min-wall-class-confidence", type=float, default=0.45)
    ap.add_argument("--min-object-confidence", type=float, default=0.5)
    ap.add_argument("--min-object-pixels", type=int, default=8)
    ap.add_argument("--window-min-area", type=int, default=64,
                    help="drop window blobs smaller than this many pixels")
    ap.add_argument("--window-merge-gap", type=int, default=8,
                    help="merge window fragments separated by up to this many pixels (mullions)")
    ap.add_argument("--window-min-fill", type=float, default=0.5,
                    help="minimum rectangle fill ratio for a window instance")
    ap.add_argument("--window-max-aspect", type=float, default=8.0,
                    help="maximum width/height ratio for a plausible window")
    ap.add_argument("--regularize-windows", action=argparse.BooleanOptionalAction, default=True,
                    help="fit window rectangles and add their keep-out to the mask")
    ap.add_argument("--keepout-margin-px", type=int, default=8,
                    help="clearance around trim, fixtures, uncertain pixels, and non-wall boundaries")
    ap.add_argument("--max-uncertain-fraction", type=float, default=0.25)
    ap.add_argument("--min-drywall-confidence", type=float, default=0.65)
    ap.add_argument("--min-wallpaper-fraction", type=float, default=0.05,
                    help="share of the wall the model must call wallpaper before reporting wallpaper_detected")
    ap.add_argument("--wallpaper-cue-threshold", type=float, default=0.5,
                    help="pattern/seam cue score that flags possible wallpaper the model did not report")
    ap.add_argument("--condition-mask", type=Path, default=Path("surface_condition_mask.png"),
                    help="0=sound, 1=paint failure, 2=rust, 255=unknown/not wall (needs a condition head)")
    ap.add_argument("--min-condition-confidence", type=float, default=None,
                    help="condition-head confidence for a damaged pixel (default: the checkpoint's calibrated "
                         f"condition_operating_point, else {DEFAULT_CONDITION[0]})")
    ap.add_argument("--min-defect-fraction", type=float, default=None,
                    help="share of the decided wall the model must call paint failure/rust to report it "
                         f"(default: the checkpoint's calibrated condition_operating_point, else {DEFAULT_CONDITION[1]})")
    ap.add_argument("--prep-margin-px", type=float, default=12.0,
                    help="how far (image pixels) a damage region's prep zone feathers past the visible damage")
    ap.add_argument("--rust-cue-threshold", type=float, default=0.5,
                    help="model-free rust cue score that flags possible rust the model did not report")
    ap.add_argument("--depth", type=Path, help="aligned depth map (.npy metres or 16-bit PNG): stop lines and area "
                                              "from fitted 3D planes (mm-level)")
    ap.add_argument("--depth-unit", choices=("m", "mm"), default="m", help="unit of a PNG depth map")
    ap.add_argument("--intrinsics", help="fx,fy,cx,cy of the RGB frame (default: EXIF focal, centred)")
    ap.add_argument("--wall-height-ft", type=float, help="measured floor-to-ceiling height (best scale input)")
    ap.add_argument("--wall-width-ft", type=float, help="measured width of the wall face (single-face frames)")
    ap.add_argument("--mm-per-px", type=float, help="depth/LiDAR scale at the wall (depth_validation report)")
    ap.add_argument("--ceiling-height-ft", type=float, default=8.0,
                    help="assumed floor-to-ceiling height when nothing is measured (reported as an assumption)")
    ap.add_argument("--focal-px", type=float, help="camera focal length in pixels (default: EXIF, else the wall itself)")
    ap.add_argument("--coats", type=int, default=2, help="finish coats for the paint estimate")
    ap.add_argument("--surface", choices=("smooth", "rough"), default="smooth",
                    help="rough = masonry or heavy texture (lower spread rate)")
    ap.add_argument("--application", choices=("roller", "brush", "spray"), default="roller")
    ap.add_argument("--waste", type=float, default=0.10, help="extra paint for waste/touch-up (fraction)")
    ap.add_argument("--no-measure", dest="measure", action="store_false",
                    help="skip stop lines, area and paint estimate")
    ap.add_argument("--tta", action="store_true", help="average normal and horizontally flipped predictions")
    ap.add_argument("--colors", action=argparse.BooleanOptionalAction, default=True,
                    help="report the paint colour of the wall: name, palette, per-face, reference match (paint_color.py)")
    ap.add_argument("--reference-color",
                    help="specified colour, checked against the coat and against the tins on site: "
                         "\"#RRGGBB\", \"r,g,b\", or a name such as \"magnolia\"")
    ap.add_argument("--color-tolerance", type=float, default=MATCH_DELTA_E,
                    help=f"largest CIEDE2000 difference that still counts as the same colour (default {MATCH_DELTA_E})")
    ap.add_argument("--color-palette-k", type=int, default=3,
                    help="how many colours to look for on the painted wall (a second one means two coats or a wrong face)")
    ap.add_argument("--objects", action=argparse.BooleanOptionalAction, default=True,
                    help="look for the equipment on site: paint cans, trays, brushes, rollers, ladders, dust "
                         "sheets (site_objects.py); also reads the colour off a tin's label")
    ap.add_argument("--min-equipment-confidence", type=float, default=DEFAULT_EQUIPMENT_CONFIDENCE,
                    help=f"cue score a piece of equipment needs (default {DEFAULT_EQUIPMENT_CONFIDENCE})")
    ap.add_argument("--required-tools", default="roller,trim",
                    help="plan tools that must be on site (default roller,trim -> a roller and a brush)")
    ap.add_argument("--objects-in-keepout", action=argparse.BooleanOptionalAction, default=True,
                    help="keep equipment found inside the paint area out of the paint plan")
    return ap


def main():
    ap = build_parser()
    args = ap.parse_args()
    for name in ("min_wall_probability", "min_wall_class_confidence", "min_object_confidence", "max_uncertain_fraction",
                 "min_drywall_confidence", "window_min_fill", "min_wallpaper_fraction", "wallpaper_cue_threshold",
                 "min_condition_confidence", "min_defect_fraction", "rust_cue_threshold",
                 "min_equipment_confidence"):
        if getattr(args, name) is not None and not 0.0 <= getattr(args, name) <= 1.0:
            ap.error(f"--{name.replace('_', '-')} must be between 0 and 1")
    if args.min_object_pixels < 1 or args.min_touchup_pixels < 1:
        ap.error("--min-object-pixels and --min-touchup-pixels must be at least 1")
    if args.coats < 1 or not 0 <= args.waste <= 1:
        ap.error("--coats must be at least 1 and --waste between 0 and 1")
    for name in ("wall_height_ft", "wall_width_ft", "mm_per_px", "focal_px", "ceiling_height_ft"):
        if getattr(args, name) is not None and getattr(args, name) <= 0:
            ap.error(f"--{name.replace('_', '-')} must be positive")
    if args.keepout_margin_px < 0:
        ap.error("--keepout-margin-px cannot be negative")
    if args.color_tolerance <= 0:
        ap.error("--color-tolerance must be greater than 0")
    if not 1 <= args.color_palette_k <= 5:
        ap.error("--color-palette-k must be between 1 and 5")

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    from hub import resolve_checkpoint
    checkpoint = torch.load(resolve_checkpoint(args.checkpoint), map_location=device, weights_only=True)
    if tuple(checkpoint["classes"]) != NAMES:
        ap.error("checkpoint class schema is incompatible; retrain with the wall/skirting labels documented in README.md")
    try:
        material_names = checkpoint_material_classes(checkpoint)
    except ValueError as exc:
        ap.error(str(exc))
    model = build_from_checkpoint(checkpoint, len(NAMES))
    model.to(device).eval()

    bgr = cv2.imread(str(args.image))
    if bgr is None:
        ap.error(f"could not read image: {args.image}")
    if args.focal_px is None:
        args.focal_px = exif_focal_px(args.image, bgr.shape[1], bgr.shape[0])
    args.depth_map = None
    if args.intrinsics:
        try:
            args.intrinsics = tuple(float(v) for v in args.intrinsics.split(","))
            assert len(args.intrinsics) == 4
        except (ValueError, AssertionError):
            ap.error("--intrinsics must be four numbers: fx,fy,cx,cy")
    if args.depth:
        from depth_validation import load_depth
        args.depth_map = load_depth(args.depth, args.depth_unit)
    summary, masks = analyze(bgr, model, checkpoint, device, args)
    h, w = bgr.shape[:2]
    (mask, wall, skirting, drywall, wallpaper, safe_paintable, hazards, uncertain, drywall_known, material_class, confident_painted, touch_up, completion) = (masks[k] for k in ("mask", "wall", "skirting", "drywall", "wallpaper", "safe_paintable", "hazards", "uncertain", "drywall_known", "material_class", "confident_painted", "touch_up", "completion"))

    overlay = cv2.addWeighted(bgr, 0.52, COLORS_BGR[mask], 0.48, 0)
    # Emphasize the predicted wall perimeter and the wall/skirting transition.
    contours, _ = cv2.findContours(wall.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    cv2.drawContours(overlay, contours, -1, (0, 255, 255), 2)
    skirt_contours, _ = cv2.findContours(skirting.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    cv2.drawContours(overlay, skirt_contours, -1, (255, 0, 255), 2)
    drywall_contours, _ = cv2.findContours(drywall.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    cv2.drawContours(overlay, drywall_contours, -1, (255, 255, 255), 1)
    wallpaper_contours, _ = cv2.findContours(wallpaper.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    cv2.drawContours(overlay, wallpaper_contours, -1, (180, 60, 255), 2)
    for region in completion["touch_up_regions"]:
        x0, y0, x1, y1 = region["bbox_xyxy"]
        cv2.rectangle(overlay, (x0, y0), (x1, y1), (0, 0, 255), 2)
    defect_contours, _ = cv2.findContours(masks["defect"].astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    cv2.drawContours(overlay, defect_contours, -1, (0, 120, 255), 2)   # orange: prepare before painting
    draw_measurements(overlay, summary.get("measurements") or {})
    draw_site_objects(overlay, summary)
    args.overlay.parent.mkdir(parents=True, exist_ok=True)
    args.mask.parent.mkdir(parents=True, exist_ok=True)
    args.paintable_mask.parent.mkdir(parents=True, exist_ok=True)
    args.keepout_mask.parent.mkdir(parents=True, exist_ok=True)
    args.uncertainty_mask.parent.mkdir(parents=True, exist_ok=True)
    args.drywall_mask.parent.mkdir(parents=True, exist_ok=True)
    args.painted_mask.parent.mkdir(parents=True, exist_ok=True)
    args.touchup_mask.parent.mkdir(parents=True, exist_ok=True)
    args.condition_mask.parent.mkdir(parents=True, exist_ok=True)
    cv2.imwrite(str(args.overlay), overlay)
    cv2.imwrite(str(args.mask), mask)
    cv2.imwrite(str(args.paintable_mask), (safe_paintable.astype(np.uint8) * 255))
    cv2.imwrite(str(args.keepout_mask), (hazards.astype(np.uint8) * 255))
    cv2.imwrite(str(args.uncertainty_mask), (uncertain.astype(np.uint8) * 255))
    cv2.imwrite(str(args.painted_mask), (confident_painted.astype(np.uint8) * 255))
    cv2.imwrite(str(args.touchup_mask), (touch_up.astype(np.uint8) * 255))
    condition_labels = np.full((h, w), 255, dtype=np.uint8)
    if masks["condition_class"] is not None:
        condition_labels[masks["known_wall"]] = masks["condition_class"][masks["known_wall"]]
    cv2.imwrite(str(args.condition_mask), condition_labels)
    drywall_labels = np.full((h, w), 255, dtype=np.uint8)
    drywall_labels[wall & drywall_known & (material_class == 0)] = 0
    drywall_labels[drywall] = 1
    drywall_labels[wallpaper] = MATERIAL_CLASSES.index("wallpaper")
    cv2.imwrite(str(args.drywall_mask), drywall_labels)
    rendered = json.dumps(summary, indent=2)
    print(rendered)
    print(f"overlay: {args.overlay}\nmask: {args.mask}\nsafe paintable mask: {args.paintable_mask}\nkeep-out mask: {args.keepout_mask}\nuncertainty mask: {args.uncertainty_mask}\ndrywall material mask: {args.drywall_mask}\n"
          f"confident painted mask: {args.painted_mask}\ntouch-up mask: {args.touchup_mask}\n"
          f"surface condition mask: {args.condition_mask}")
    if args.json:
        args.json.parent.mkdir(parents=True, exist_ok=True)
        args.json.write_text(rendered + "\n")


if __name__ == "__main__":
    main()