#!/usr/bin/env python3 """Turn segmentation masks into an executable robot paint plan. The semantic model answers "what is wall"; a robot needs "where do I move". This module closes that gap. Given the paintable and keep-out masks that ``predict.py`` writes (or a raw semantic mask), it produces: * the wall's outer boundary polygon and its corners, in image coordinates and, when four wall corners are supplied, in rectified wall-plane coordinates; * per-edge lengths in plane units and millimetres; * a two-tool stroke plan: roller passes for open areas plus optional trim/brush passes (``--trim-width``) for edges and around fixtures the roller cannot reach, with configurable pass width and overlap; * machine-ready waypoints with paint on/off flags and travel moves; * paint-progress tracking: an observation grid (``--state``) records what has been painted, the plan targets only what remains, and the updated grid is written back so the robot knows when the wall is done. Geometry uses numpy + scipy only, so the planner runs without torch/cv2 and is easy to unit-test. Coordinates are image pixels unless ``--corners`` rectifies the wall; supply ``--mm-per-unit`` to also report millimetres. Example:: python3 robot_planner.py --mask painting_mask.png \ --corners "80,60 1180,55 1190,700 70,710" \ --stroke-width 230 --trim-width 60 --overlap 0.2 --margin 40 --mm-per-unit 3.0 \ --state wall_014.npz --plan plan.json --overlay plan.png """ from __future__ import annotations import argparse import sys import json from pathlib import Path import numpy as np from PIL import Image, ImageDraw from scipy import ndimage WALL_CLASSES = (1, 2, 3) # wall_unpainted, wall_painted, wall_uncertain KEEPOUT_CLASSES = (4, 5, 6, 7, 8, 9) # skirting, switch/outlet, AC, door, window, obstacle def disk(radius): size = int(radius) * 2 + 1 ys, xs = np.mgrid[0:size, 0:size] return (xs - radius) ** 2 + (ys - radius) ** 2 <= radius ** 2 def largest_component(mask): """Largest 8-connected blob, so stray specks do not define the wall.""" labels, count = ndimage.label(mask, structure=np.ones((3, 3), dtype=int)) if count == 0: return np.zeros_like(mask, dtype=bool) sizes = ndimage.sum(mask, labels, index=np.arange(1, count + 1)) return labels == (int(np.argmax(sizes)) + 1) MOORE = [(-1, 0), (-1, 1), (0, 1), (1, 1), (1, 0), (1, -1), (0, -1), (-1, -1)] def trace_boundary(mask): """Ordered outer boundary of a mask as (y, x) points (Moore-neighbour trace). Unlike a convex hull this follows concavities (rooflines, sky cut-outs), so the reported wall boundary is the real one. Holes are handled separately: keep-out regions never enter stroke geometry and are not part of this loop. """ height, width = mask.shape start = None for y in range(height): row = np.nonzero(mask[y])[0] if row.size: start = (y, int(row[0])) break if start is None: return [] contour = [start] backtrack = (start[0], start[1] - 1) current = start guard = 8 * height * width while guard > 0: guard -= 1 delta = (backtrack[0] - current[0], backtrack[1] - current[1]) index = MOORE.index(delta) if delta in MOORE else 0 found = None for step in range(1, 9): direction = MOORE[(index + step) % 8] ny, nx = current[0] + direction[0], current[1] + direction[1] if 0 <= ny < height and 0 <= nx < width and mask[ny, nx]: previous = MOORE[(index + step - 1) % 8] found = (ny, nx) backtrack = (current[0] + previous[0], current[1] + previous[1]) break if found is None: break current = found contour.append(current) if current == start and len(contour) > 3: break return contour def _douglas_peucker(points, tolerance): if len(points) < 3: return points start, end = np.asarray(points[0], float), np.asarray(points[-1], float) line = end - start length = np.linalg.norm(line) offsets = np.asarray(points, float) - start if length == 0: distances = np.linalg.norm(offsets, axis=1) else: distances = np.abs(line[0] * offsets[:, 1] - line[1] * offsets[:, 0]) / length index = int(np.argmax(distances)) if distances[index] <= tolerance: return [points[0], points[-1]] left = _douglas_peucker(points[:index + 1], tolerance) right = _douglas_peucker(points[index:], tolerance) return left[:-1] + right def simplify_closed(polygon, tolerance): """Douglas-Peucker on a closed polygon by splitting at the farthest point.""" if len(polygon) < 4: return polygon points = list(polygon) first = np.asarray(points[0], float) distances = [np.linalg.norm(np.asarray(p, float) - first) for p in points] split = int(np.argmax(distances)) chain_a = _douglas_peucker(points[:split + 1], tolerance) chain_b = _douglas_peucker(points[split:] + [points[0]], tolerance) simplified = chain_a[:-1] + chain_b[:-1] return simplified if len(simplified) >= 3 else list(polygon) def homography(source, destination): """3x3 homography mapping source points (N, 2) to destination points (N, 2).""" source, destination = np.asarray(source, float), np.asarray(destination, float) rows = [] for (x, y), (X, Y) in zip(source, destination): rows.append([-x, -y, -1, 0, 0, 0, x * X, y * X, X]) rows.append([0, 0, 0, -x, -y, -1, x * Y, y * Y, Y]) _, _, solution = np.linalg.svd(np.asarray(rows)) return solution[-1].reshape(3, 3) def apply_homography(matrix, points): points = np.asarray(points, float) homogeneous = np.c_[points, np.ones(len(points))] mapped = homogeneous @ matrix.T return mapped[:, :2] / mapped[:, 2:3] def warp_mask(mask, matrix, size): """Nearest-neighbour warp of a boolean mask into ``size`` = (width, height).""" width, height = size inverse = np.linalg.inv(matrix) ys, xs = np.mgrid[0:height, 0:width] source = apply_homography(inverse, np.c_[xs.ravel(), ys.ravel()]) sx, sy = np.rint(source[:, 0]).astype(int), np.rint(source[:, 1]).astype(int) valid = (sx >= 0) & (sx < mask.shape[1]) & (sy >= 0) & (sy < mask.shape[0]) output = np.zeros((height, width), dtype=bool) output.ravel()[valid] = mask[sy[valid], sx[valid]] return output def runs_of_true(row): index = np.flatnonzero(row) if index.size == 0: return [] breaks = np.where(np.diff(index) > 1)[0] starts = np.concatenate(([0], breaks + 1)) ends = np.concatenate((breaks, [index.size - 1])) return [(int(index[s]), int(index[e])) for s, e in zip(starts, ends)] def _row_runs(grid, spacing, min_run): """[(line, [(start, end), ...]), ...] top to bottom: passes anchored to the region's extent (not the image centre) so the first and last pass reach the margin instead of leaving an unplanned gap.""" height = grid.shape[0] occupied = np.flatnonzero(grid.any(axis=1)) if occupied.size == 0: return [] y_low, y_high = int(occupied[0]), int(occupied[-1]) count = max(1, int(round((y_high - y_low) / spacing)) + 1) offsets = np.linspace(y_low, y_high, count) if count > 1 else np.array([float(y_low)]) rows = [] for offset in offsets: line = int(round(offset)) if line >= height: continue runs = [(a, b) for a, b in runs_of_true(grid[line]) if b - a >= min_run] if runs: rows.append((line, runs)) return rows def _cells(rows): """Boustrophedon cellular decomposition of the passes: consecutive rows' runs that overlap one-to-one belong to one cell (a vertical band between obstacles); a run that splits around an obstacle, or joins after it, starts new cells. Returns cells as lists of (line, start, end), each ordered top to bottom.""" cells, owner_prev, prev_runs = [], {}, [] for line, runs in rows: owner = {} for k, (a, b) in enumerate(runs): above = [j for j, (pa, pb) in enumerate(prev_runs) if pa <= b and a <= pb] joined = None if len(above) == 1: j = above[0] pa, pb = prev_runs[j] below_of_j = [i for i, (na, nb) in enumerate(runs) if pa <= nb and na <= pb] if len(below_of_j) == 1: joined = owner_prev[j] if joined is None: cells.append([]) joined = len(cells) - 1 cells[joined].append((line, a, b)) owner[k] = joined owner_prev, prev_runs = owner, runs return cells def plan_strokes(allowed, spacing, orientation, min_run, order="optimal"): """Segments (roller centre lines) over the allowed region. ``spacing`` is the distance between passes; ``orientation`` is horizontal (passes run left-right) or vertical. Rows blocked by keep-out break into separate segments so the robot can lift between them. ``order="optimal"`` (default) is ``"cells"`` with the cell visiting order solved exactly (path_motion.sequence_cells: top-down precedence, either entry end, hops and lifts priced); it falls back to ``"cells"`` on very many cells. ``order="cells"`` paints one vertical band at a time, top to bottom, serpentine inside the band, so a window is passed once instead of being crossed on every row; ``order="rows"`` is the old sweep of the whole wall row by row (kept for comparison: 40-75 % of its path is travel on a wall with windows). """ grid = allowed if orientation == "horizontal" else allowed.T if spacing <= 0: raise ValueError("spacing must be positive") rows = _row_runs(grid, spacing, min_run) segments = [] if order == "rows": forward = True for line, runs in rows: for start, end in (runs if forward else runs[::-1]): segments.append([(start, line), (end, line)] if forward else [(end, line), (start, line)]) forward = not forward else: cells = _cells(rows) best = None if order == "optimal": import path_motion best = path_motion.sequence_cells(cells, spacing) if best is not None: segments = path_motion.strokes_in_order(cells, best) cells = [] # top of each cell first, then left to right; a cell's first stroke starts at the end nearest the robot todo = sorted(range(len(cells)), key=lambda c: (cells[c][0][0], cells[c][0][1])) here = None while todo: if here is None: pick = todo[0] else: # among the cells whose top is no lower than the highest remaining top + a pass, the nearest top = cells[todo[0]][0][0] near = [c for c in todo if cells[c][0][0] <= top + spacing * 1.5] or todo[:1] pick = min(near, key=lambda c: min(abs(cells[c][0][1] - here[0]), abs(cells[c][0][2] - here[0])) + abs(cells[c][0][0] - here[1])) todo.remove(pick) cell = cells[pick] first = cell[0] forward = here is None or abs(first[1] - here[0]) <= abs(first[2] - here[0]) for line, start, end in cell: segments.append([(start, line), (end, line)] if forward else [(end, line), (start, line)]) forward = not forward here = segments[-1][-1] if orientation == "vertical": segments = [[(y, x) for x, y in segment] for segment in segments] return segments def footprint(segments, shape, width, exact=False): """Approximate painted footprint of centre-line segments as thick bands. ``exact=True`` is a roller's real contact: a swath ``width`` wide ACROSS the stroke that ends where the stroke ends (the square, half-a-width-longer-at-each-end footprint over-claims the ends of every pass).""" mask = np.zeros(shape, dtype=bool) half = width / 2.0 for (x0, y0), (x1, y1) in segments: along_x = abs(x1 - x0) >= abs(y1 - y0) hx, hy = (0.0, half) if exact and along_x else (half, 0.0) if exact else (half, half) y_low, y_high = int(np.floor(min(y0, y1) - hy)), int(np.ceil(max(y0, y1) + hy)) x_low, x_high = int(np.floor(min(x0, x1) - hx)), int(np.ceil(max(x0, x1) + hx)) y_low, x_low = max(0, y_low), max(0, x_low) y_high, x_high = min(shape[0], y_high + 1), min(shape[1], x_high + 1) mask[y_low:y_high, x_low:x_high] = True return mask def path_length(segments): total = 0.0 previous = None for segment in segments: start, end = np.asarray(segment[0], float), np.asarray(segment[-1], float) total += float(np.linalg.norm(end - start)) if previous is not None: total += float(np.linalg.norm(start - previous)) previous = end return total def travel_distance(segments): """Non-painting travel between segments (path length minus painted strokes).""" painted = sum(float(np.linalg.norm(np.asarray(s[-1], float) - np.asarray(s[0], float))) for s in segments) return path_length(segments) - painted def plan_pass(region, wall_distance, keepout_distance, width, overlap, margin, orientation, min_run, order="optimal", edge_reach=0.0): """Plan one tool pass over ``region``. The centre line must stay ``margin + width/2`` inside the wall edge (so the footprint never leaves the wall) and ``width/2`` off fixtures (the tool edge may meet a frame). ``edge_reach`` lets the footprint run that far past those limits, for edges that are taped (the roller may touch the tape). Returns the centre-line segments and the allowed mask. """ inset = margin + width / 2.0 allowed = region & (wall_distance >= inset - edge_reach) & (keepout_distance >= width / 2.0 - edge_reach) spacing = max(1.0, width * (1.0 - overlap)) return plan_strokes(allowed, spacing, orientation, min_run=min_run, order=order), allowed, spacing def extend_strokes(segments, reachable, width, orientation): """Run each roller pass on to the wall edge: the centre line may go on while the whole swath across it still lies in ``reachable`` (the target plus the taped margin). The planner keeps centres half a width from the edge in both directions, which is right across a pass but leaves a strip at every pass end.""" half = int(np.floor(width / 2.0)) out = [] h, w = reachable.shape for seg in segments: (x0, y0), (x1, y1) = seg[0], seg[-1] horizontal = abs(x1 - x0) >= abs(y1 - y0) grid = reachable if horizontal else reachable.T line = int(round(y0 if horizontal else x0)) a, b = (x0, x1) if horizontal else (y0, y1) sign = 1 if b >= a else -1 def room(p): lo, hi = max(0, line - half), min(grid.shape[0], line + half + 1) return 0 <= p < grid.shape[1] and grid[lo:hi, p].all() for end, direction in ((b, sign), (a, -sign)): p = int(round(end)) while room(p + direction): p += direction if end is b: b = p else: a = p out.append([(a, y0), (b, y1)] if horizontal else [(x0, a), (x1, b)]) return out def plan_trim_pieces(region, wall_distance, keepout_distance, width, overlap, order="optimal", edge_reach=0.0): """Brush strokes for what the roller left, piece by piece, each along its own long axis (a strip beside a door frame is brushed up and down, not in many short sideways dabs). Pieces go top to bottom, then left to right.""" labels, count = ndimage.label(region) spacing = max(1.0, width * (1.0 - overlap)) pieces = [] for k, box in enumerate(ndimage.find_objects(labels), start=1): if box is None: continue height, wide = box[0].stop - box[0].start, box[1].stop - box[1].start orientation = "vertical" if height > wide * 1.25 else "horizontal" piece = labels == k segments, _, _ = plan_pass(piece, wall_distance, keepout_distance, width, overlap, 0.0, orientation, max(1, width * 0.5) if not edge_reach else 1, order, edge_reach) if segments: pieces.append((box[0].start, box[1].start, segments)) pieces.sort(key=lambda p: (p[0] // max(1.0, width * 4), p[1])) return [segment for _, _, segments in pieces for segment in segments], spacing def choose_orientation(region, wall_distance, keepout_distance, width, overlap, margin, order, mm_per_unit, speed_mm_s=250.0, horizontal_bias=0.05): """"horizontal" or "vertical", whichever the timeline says is faster (hops and lifts included). Horizontal is the habit (top-down, drips fall on bare wall) so vertical must win by ``horizontal_bias`` to be chosen.""" import path_motion seconds = {} for orientation in ("horizontal", "vertical"): segments, _, _ = plan_pass(region, wall_distance, keepout_distance, width, overlap, margin, orientation, width, order) seconds[orientation] = path_motion.timeline(segments, mm_per_unit, speed_mm_s, width_units=width)[1] \ if segments else float("inf") return "vertical" if seconds["vertical"] < seconds["horizontal"] * (1.0 - horizontal_bias) else "horizontal" def sequence_waypoints(segments_by_tool): """Machine-ready waypoint order with travel moves and paint flags. Each segment contributes a paint-on point and a paint-off point; a travel point (``paint: false``) is inserted before each segment after the first so a controller can command gun on/off and roller lift. """ waypoints = [] for tool, segments in segments_by_tool: for index, segment in enumerate(segments): start = segment[0] end = segment[-1] if waypoints: waypoints.append({"x": round(float(start[0]), 2), "y": round(float(start[1]), 2), "paint": False, "tool": "travel"}) waypoints.append({"x": round(float(start[0]), 2), "y": round(float(start[1]), 2), "paint": True, "tool": tool, "segment": index}) waypoints.append({"x": round(float(end[0]), 2), "y": round(float(end[1]), 2), "paint": True, "tool": tool, "segment": index}) return waypoints def load_state(path): """Load a persisted paint-progress grid as (painted, plane_size) or None.""" path = Path(path) if not path.is_file(): return None with np.load(path) as state: painted = state["painted"].astype(bool) size = tuple(int(v) for v in state["size"]) return painted, size def save_state(path, painted, size): path = Path(path) path.parent.mkdir(parents=True, exist_ok=True) np.savez_compressed(path, painted=painted.astype(np.uint8), size=np.asarray(size, dtype=int)) def count_passes(segments, orientation): axis = 1 if orientation == "horizontal" else 0 return len({round(segment[0][axis], 3) for segment in segments}) def build_plan(paintable, keepout, stroke_width, overlap, margin, orientation, corners, mm_per_unit, boundary_tolerance, trim_width=None, trim_overlap=None, already_painted=None, plane_size=None, order="optimal", edge_reach=0.0, trim_orientation="auto", exact_ends=False): """``exact_ends=True`` plans a roller's true contact: passes run on to the edge, the footprint ends where the pass ends, the brush may sit right on a taped edge, every wall piece is painted (not only the largest), and a last sweep brushes anything still bare. ``edge_reach`` (units): the roller may overlap taped edges by this much, so less is left for the brush. ``trim_orientation="auto"`` runs the brush along each leftover strip's long axis (horizontal/vertical per piece) instead of across it. ``plane_size`` (units wide, units high) sets the true wall-plane aspect; without it the longest image edges are used, which is only right for a near-frontal view.""" height, width = paintable.shape given_size, plane_size, to_plane, to_image = plane_size, (width, height), None, None if corners is not None: corners = np.asarray(corners, float) edge_w = max(np.linalg.norm(corners[1] - corners[0]), np.linalg.norm(corners[2] - corners[3])) edge_h = max(np.linalg.norm(corners[3] - corners[0]), np.linalg.norm(corners[2] - corners[1])) plane_size = given_size or (max(1, int(round(edge_w))), max(1, int(round(edge_h)))) to_plane = homography(corners, [(0, 0), (plane_size[0], 0), (plane_size[0], plane_size[1]), (0, plane_size[1])]) to_image = np.linalg.inv(to_plane) paintable = warp_mask(paintable, to_plane, plane_size) keepout = warp_mask(keepout, to_plane, plane_size) if exact_ends: labels, count = ndimage.label(paintable) sizes = ndimage.sum(paintable, labels, range(1, count + 1)) if count else [] keep = [i + 1 for i, a in enumerate(sizes) if a >= max(40, 0.002 * paintable.size)] paintable = np.isin(labels, keep) if keep else paintable else: paintable = largest_component(paintable) if not paintable.any(): raise ValueError("paintable mask is empty after the largest-component filter") target = paintable & ~keepout if already_painted is not None and already_painted.shape != target.shape: raise ValueError(f"paint-progress grid {already_painted.shape} does not match plane {target.shape}") already = (already_painted & target) if already_painted is not None else np.zeros_like(target) pending = target & ~already polygon = simplify_closed([(x, y) for y, x in trace_boundary(paintable)], boundary_tolerance) wall_distance = ndimage.distance_transform_edt(paintable) keepout_distance = ndimage.distance_transform_edt(~keepout) if orientation == "auto": orientation = choose_orientation(pending, wall_distance, keepout_distance, stroke_width, overlap, margin, order, mm_per_unit or 1.0) roller_segments, roller_allowed, spacing = plan_pass( pending, wall_distance, keepout_distance, stroke_width, overlap, margin, orientation, stroke_width, order, edge_reach) if exact_ends: reachable = ndimage.binary_dilation(target, iterations=max(1, int(round(edge_reach)))) if edge_reach else target roller_segments = extend_strokes(roller_segments, reachable, stroke_width, orientation) roller_footprint = footprint(roller_segments, target.shape, stroke_width, exact=exact_ends) & target trim_segments, trim_spacing = [], None trim_footprint = np.zeros_like(target) if trim_width and trim_width > 0: remaining = pending & ~roller_footprint t_overlap = overlap if trim_overlap is None else trim_overlap if trim_orientation == "auto": trim_segments, trim_spacing = plan_trim_pieces(remaining, wall_distance, keepout_distance, trim_width, t_overlap, order, trim_width / 2.0 if exact_ends else 0.0) else: trim_segments, _, trim_spacing = plan_pass( remaining, wall_distance, keepout_distance, trim_width, t_overlap, 0.0, orientation, max(1, trim_width * 0.5), order) trim_footprint = footprint(trim_segments, target.shape, trim_width) & target if exact_ends: # last sweep: whatever is still bare, brushed wherever the brush can sit bare = pending & ~roller_footprint & ~trim_footprint if bare.any(): big = np.full(target.shape, 1e9) extra, _ = plan_trim_pieces(bare, big, big, trim_width, t_overlap, order, trim_width / 2.0) trim_segments = trim_segments + extra trim_footprint = footprint(trim_segments, target.shape, trim_width) & target newly = (roller_footprint | trim_footprint) & target target_area = int(target.sum()) covered_pending = int((newly & pending).sum()) covered_total = int(((already | newly) & target).sum()) planned_fraction = covered_pending / target_area if target_area else 0.0 residual_fraction = 1.0 - (covered_total / target_area) if target_area else 0.0 pending_area = int(pending.sum()) slack = np.minimum(wall_distance - (margin + stroke_width / 2.0), keepout_distance - stroke_width / 2.0) max_inradius = float(slack[target].max()) if target.any() else 0.0 warnings = [] if target_area == 0: warnings.append("paintable region is empty after excluding keep-out") elif pending_area == 0: warnings.append("target already fully painted according to the progress grid") elif not (roller_segments or trim_segments): warnings.append(f"no tool centre-line space: inset {margin + stroke_width / 2.0:g} exceeds the wall's " f"inradius ({max_inradius:.1f}); reduce --stroke-width or --margin") elif residual_fraction > 0.2 and not trim_segments: warnings.append(f"roller plan leaves {residual_fraction:.0%} unpainted; add --trim-width for " f"edges and around fixtures") def transform(points): return apply_homography(to_plane, points) if to_plane is not None else np.asarray(points, float) polygon_plane = transform(polygon) edges = [] for i in range(len(polygon_plane)): start, end = polygon_plane[i], polygon_plane[(i + 1) % len(polygon_plane)] length = float(np.linalg.norm(np.asarray(end) - np.asarray(start))) edges.append({"from": [round(float(v), 2) for v in start], "to": [round(float(v), 2) for v in end], "length_units": round(length, 2), "length_mm": round(length * mm_per_unit, 2) if mm_per_unit else None}) scale = mm_per_unit or 1.0 tools = [name for name, segs in (("roller", roller_segments), ("trim", trim_segments)) if segs] waypoints = sequence_waypoints([("roller", roller_segments), ("trim", trim_segments)]) all_segments = roller_segments + trim_segments def strokes(segments): return [[[round(float(x), 2), round(float(y), 2)] for x, y in segment] for segment in segments] return { "units": "mm" if mm_per_unit else "plane-pixels", "rectified": to_plane is not None, "warnings": warnings, "plane_size": {"width": plane_size[0], "height": plane_size[1], "mm_per_unit": mm_per_unit}, "wall": { "boundary_polygon_plane": [[round(float(x), 2), round(float(y), 2)] for x, y in polygon_plane], "boundary_polygon_image": ([[round(float(x), 2), round(float(y), 2)] for x, y in apply_homography(to_image, polygon_plane)] if to_image is not None else None), "corners": len(polygon_plane), "edges": edges, "area_units2": round(target_area * (scale ** 2), 2), "area_mm2": round(target_area * (scale ** 2), 2) if mm_per_unit else None, }, "paint_region": {"target_area_units2": round(target_area * (scale ** 2), 2), "pending_area_units2": round(pending_area * (scale ** 2), 2), "roller_inset_units": round(margin + stroke_width / 2.0, 2)}, "tools": tools, "plan": { "orientation": orientation, "stroke_width_units": stroke_width, "overlap": overlap, "spacing_units": round(spacing, 2), "margin_units": margin, "passes": count_passes(roller_segments, orientation), "segments": len(roller_segments), "path_length_units": round(path_length(roller_segments), 2), "path_length_mm": round(path_length(roller_segments) * mm_per_unit, 2) if mm_per_unit else None, "strokes": strokes(roller_segments), "trim": { "enabled": bool(trim_segments), "stroke_width_units": trim_width, "spacing_units": round(trim_spacing, 2) if trim_spacing else None, "passes": count_passes(trim_segments, orientation), "segments": len(trim_segments), "path_length_units": round(path_length(trim_segments), 2), "strokes": strokes(trim_segments), }, "waypoints": waypoints, "travel_distance_units": round(travel_distance(all_segments), 2), "travel_distance_mm": round(travel_distance(all_segments) * mm_per_unit, 2) if mm_per_unit else None, }, "coverage": { "already_painted_fraction": round(int(already.sum()) / target_area, 4) if target_area else 0.0, "planned_fraction": round(planned_fraction, 4), "total_covered_fraction": round(covered_total / target_area, 4) if target_area else 0.0, "remaining_fraction": round(residual_fraction, 4), "uncovered_fraction": round(residual_fraction, 4), "planned_area_units2": round(covered_pending * (scale ** 2), 2), "remaining_area_units2": round((target_area - covered_total) * (scale ** 2), 2), "remaining_area_mm2": round((target_area - covered_total) * (scale ** 2), 2) if mm_per_unit else None, }, }, {"target": target, "allowed": roller_allowed, "roller": roller_segments, "trim": trim_segments, "painted": (already | newly) & target, "polygon": polygon_plane, "size": plane_size, "keepout": keepout} def render_overlay(paintable, keepout, debug, path): width, height = debug["size"] canvas = np.zeros((height, width, 3), dtype=np.uint8) canvas[debug["target"]] = (70, 70, 70) canvas[debug["allowed"]] = (0, 60, 0) keepout = debug.get("keepout", keepout) if keepout.shape == debug["target"].shape: canvas[keepout] = (120, 0, 0) image = Image.fromarray(canvas) draw = ImageDraw.Draw(image) for segment in debug["trim"]: draw.line([tuple(segment[0]), tuple(segment[-1])], fill=(255, 140, 0), width=2) for segment in debug["roller"]: draw.line([tuple(segment[0]), tuple(segment[-1])], fill=(0, 160, 255), width=2) if len(debug["polygon"]) >= 2: draw.line([tuple(p) for p in debug["polygon"]] + [tuple(debug["polygon"][0])], fill=(255, 215, 0), width=2) Path(path).parent.mkdir(parents=True, exist_ok=True) image.save(path) def load_binary(path): return np.asarray(Image.open(path).convert("L")) > 0 def corners_arg(value): try: points = [[float(v) for v in pair.split(",")] for pair in value.split()] if len(points) != 4 or any(len(point) != 2 for point in points): raise ValueError return points except ValueError as exc: raise argparse.ArgumentTypeError( "expected four x,y pairs (top-left, top-right, bottom-right, bottom-left)") from exc def main(): parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) source = parser.add_mutually_exclusive_group(required=True) source.add_argument("--paintable", type=Path, help="binary paintable-wall mask PNG from predict.py") source.add_argument("--mask", type=Path, help="semantic mask PNG; paintable/keep-out are derived from it") parser.add_argument("--keepout", type=Path, help="binary keep-out mask PNG (used with --paintable)") parser.add_argument("--measurements", type=Path, help="predict.py --json output: plan inside a measured face's paint region at true size; " "plane units are then --unit-mm millimetres and widths/margins are in mm") parser.add_argument("--face", type=int, default=0, help="which measured face to plan (with --measurements)") parser.add_argument("--unit-mm", type=float, default=10.0, help="plane resolution with --measurements") parser.add_argument("--corners", type=corners_arg, help='four wall corners "TL TR BR BL" as "x,y x,y x,y x,y"') parser.add_argument("--stroke-width", type=float, default=230.0, help="roller/pass width in plane units") parser.add_argument("--overlap", type=float, default=0.2, help="fractional overlap between passes (0-0.9)") parser.add_argument("--trim-width", type=float, help="brush/trim tool width for edges and fixtures the roller cannot reach") parser.add_argument("--trim-overlap", type=float, help="overlap for the trim tool (defaults to --overlap)") parser.add_argument("--margin", type=float, default=40.0, help="clearance from wall edge and fixtures") parser.add_argument("--orientation", choices=("horizontal", "vertical"), default="horizontal") parser.add_argument("--mm-per-unit", type=float, help="physical scale for millimetre reporting") parser.add_argument("--boundary-tolerance", type=float, default=8.0, help="polygon simplification tolerance") parser.add_argument("--state", type=Path, help="paint-progress grid (.npz): loaded if present and saved after planning") parser.add_argument("--observe", type=Path, help="binary plane-aligned mask of newly observed painted pixels to merge into --state") parser.add_argument("--plan", type=Path, help="write the plan JSON here") parser.add_argument("--overlay", type=Path, help="write a visualization PNG here") args = parser.parse_args() if not 0.0 <= args.overlap < 1.0: parser.error("--overlap must be in [0, 1)") if args.trim_width is not None and args.trim_width <= 0: parser.error("--trim-width must be positive") if args.trim_overlap is not None and not 0.0 <= args.trim_overlap < 1.0: parser.error("--trim-overlap must be in [0, 1)") if args.stroke_width <= 0 or args.margin < 0: parser.error("--stroke-width must be positive and --margin non-negative") if args.paintable: paintable = load_binary(args.paintable) keepout = load_binary(args.keepout) if args.keepout else np.zeros_like(paintable) else: semantic = np.asarray(Image.open(args.mask)) if semantic.ndim == 3: parser.error("--mask expects a single-channel indexed PNG; use --paintable for a binary mask") paintable = np.isin(semantic, WALL_CLASSES) keepout = np.isin(semantic, KEEPOUT_CLASSES) plane_size = None if args.measurements: data = json.loads(args.measurements.read_text(encoding="utf-8")) faces = (data.get("measurements") or data).get("faces") or [] if not 0 <= args.face < len(faces): parser.error(f"{args.measurements} has {len(faces)} measured face(s); --face {args.face} is not one") face = faces[args.face] open_ends = [side for side, kind in face["ends"].items() if kind == "frame"] if open_ends: print(f"warning: the wall continues past the image on the {' and '.join(open_ends)}; " f"this plan covers only the visible part and has no stop there", file=sys.stderr) doubtful = [k for k, v in face.get("stop_line_trust", {}).items() if v == "conflict"] if doubtful: print(f"warning: {', '.join(doubtful)} stop line(s) disagree with the segmentation; verify before painting", file=sys.stderr) # Plan in true wall-plane millimetres: ceiling line to top of skirting, end to end. args.corners = [tuple(p) for p in face["paint_region_corners_px"]] args.mm_per_unit = args.unit_mm plane_size = (max(1, round(face["width_ft"] * 304.8 / args.unit_mm)), max(1, round(face["paint_height_ft"] * 304.8 / args.unit_mm))) args.stroke_width /= args.unit_mm args.margin /= args.unit_mm if args.trim_width: args.trim_width /= args.unit_mm already = None if args.state: loaded = load_state(args.state) if loaded is not None: painted_grid, size = loaded if size != (paintable.shape[1], paintable.shape[0]): parser.error(f"--state grid is {size}, mask is {(paintable.shape[1], paintable.shape[0])}; " f"use the same --corners/resolution") already = painted_grid if args.observe: observed = load_binary(args.observe) if observed.shape != paintable.shape: parser.error(f"--observe mask {observed.shape} does not match the input {paintable.shape}") already = observed if already is None else (already | observed) plan, debug = build_plan(paintable, keepout, args.stroke_width, args.overlap, args.margin, args.orientation, args.corners, args.mm_per_unit, args.boundary_tolerance, plane_size=plane_size, trim_width=args.trim_width, trim_overlap=args.trim_overlap, already_painted=already) rendered = json.dumps(plan, indent=2) print(rendered) if args.plan: args.plan.parent.mkdir(parents=True, exist_ok=True) args.plan.write_text(rendered + "\n", encoding="utf-8") if args.state: save_state(args.state, debug["painted"], debug["size"]) if args.overlay: render_overlay(paintable, keepout, debug, args.overlay) print(f"overlay: {args.overlay}") if __name__ == "__main__": main()