robot_planner.py · constructelligence/painting-vision-robotics-kit at main

painting-vision-robotics-kit / robot_planner.py
Update code: real-machine G-code adapter, commissioning gate, machine_job runner, pump watchdog firmware, twin completion gate, realism and masking modules (card and weights unchanged)
4eecd03 verified 3 days ago
Raw History Blame Contribute Delete
37.9 kB
#!/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()