Upload robot_planner.py with huggingface_hub

robot_planner.py

@@ -30,6 +30,7 @@ Example::
 from __future__ import annotations
 
 import argparse
+import sys
 import json
 from pathlib import Path
 
@@ -303,14 +304,16 @@ def count_passes(segments, orientation):
 
 
 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):
+               boundary_tolerance, trim_width=None, trim_overlap=None, already_painted=None, plane_size=None):
+    """``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
-    plane_size, to_plane, to_image = (width, height), None, None
+    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 = (max(1, int(round(edge_w))), max(1, int(round(edge_h))))
+        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)
@@ -484,6 +487,11 @@ def main():
     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")
@@ -522,6 +530,31 @@ def main():
         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)
@@ -539,6 +572,7 @@ def main():
 
     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)