#!/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()