progress

tools/mask_to_polygon.py

@@ -41,16 +41,16 @@ def find_largest_contour(contours: list) -> np.ndarray:
     return max(contours, key=cv2.contourArea)
 
 
-def contours_to_polygon(contours: list, mode: str) -> np.ndarray:
+def contours_to_polygon(contours: list, mode: str) -> np.ndarray | None:
     """
     Convert contours to a single polygon.
 
     Args:
         contours: List of contours from cv2.findContours
-        mode: "convex_hull" or "largest_only"
+        mode: "convex_hull", "largest_only", or "multiple"
 
     Returns:
-        Single contour as numpy array
+        Single contour as numpy array (or None for "multiple" mode)
     """
     if not contours:
         return None
@@ -62,9 +62,23 @@ def contours_to_polygon(contours: list, mode: str) -> np.ndarray:
         all_points = np.vstack(contours)
         return cv2.convexHull(all_points)
 
+    if mode == "multiple":
+        return None
+
     raise ValueError(f"Unknown mode: {mode}")
 
 
+def process_contour(
+    contour: np.ndarray, padding: float, target_points: int
+) -> np.ndarray:
+    """
+    Process a single contour: pad, simplify, and ensure clockwise.
+    """
+    padded = pad_polygon(contour, padding)
+    simplified = simplify_to_target_points(padded, target_points)
+    return ensure_clockwise(simplified)
+
+
 def simplify_to_target_points(contour: np.ndarray, target_points: int) -> np.ndarray:
     """
     Simplify a contour to approximately target number of points.
@@ -230,9 +244,9 @@ def main():
     )
     parser.add_argument(
         "--mode",
-        choices=["convex_hull", "largest_only"],
+        choices=["convex_hull", "largest_only", "multiple"],
         default="convex_hull",
-        help="How to handle multiple regions (default: convex_hull)",
+        help="How to handle multiple regions: convex_hull, largest_only, or multiple (default: convex_hull)",
     )
     parser.add_argument(
         "--preview",
@@ -240,6 +254,12 @@ def main():
         default=None,
         help="Output preview PNG showing polygon on mask",
     )
+    parser.add_argument(
+        "--min-area",
+        type=int,
+        default=100,
+        help="Minimum contour area to include in multiple mode (default: 100)",
+    )
 
     args = parser.parse_args()
 
@@ -259,25 +279,82 @@ def main():
         print("Error: No contours found in mask", file=sys.stderr)
         sys.exit(1)
 
-    polygon = contours_to_polygon(contours, args.mode)
-    padded = pad_polygon(polygon, args.padding)
-    simplified = simplify_to_target_points(padded, args.target_points)
-    clockwise = ensure_clockwise(simplified)
+    if args.mode == "multiple":
+        contours = sorted(contours, key=cv2.contourArea, reverse=True)
+        contours = [c for c in contours if cv2.contourArea(c) >= args.min_area]
 
-    output_path = args.output
-    if output_path is None:
-        output_path = args.image.with_suffix(".tres")
+        if not contours:
+            print("Error: No contours meet minimum area requirement", file=sys.stderr)
+            sys.exit(1)
 
-    uid = generate_godot_uid()
-    write_tres(clockwise, output_path, uid)
+        output_base = args.output if args.output else args.image.with_suffix("")
+        output_dir = output_base.parent
+        output_stem = output_base.stem
 
-    if args.preview:
-        write_preview(image, clockwise, args.preview)
-        print(f"Preview: {args.preview}")
+        all_points = []
+        for i, contour in enumerate(contours):
+            processed = process_contour(contour, args.padding, args.target_points)
+            all_points.append(processed)
 
-    print(f"Created: {output_path}")
-    print(f"Points: {len(clockwise.squeeze())}")
-    print(f"UID: uid://{uid}")
+            output_path = output_dir / f"{output_stem}_{i}.tres"
+            uid = generate_godot_uid()
+            write_tres(processed, output_path, uid)
+            print(f"Created: {output_path}")
+            print(f"  Points: {len(processed.squeeze())}")
+            print(f"  UID: uid://{uid}")
+
+        if args.preview:
+            preview = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
+            colors = [
+                (0, 255, 0),
+                (255, 0, 0),
+                (0, 0, 255),
+                (255, 255, 0),
+                (255, 0, 255),
+                (0, 255, 255),
+            ]
+            for i, points in enumerate(all_points):
+                color = colors[i % len(colors)]
+                pts = points.squeeze().astype(np.int32).reshape(-1, 1, 2)
+                cv2.polylines(preview, [pts], True, color, 2)
+                for j, pt in enumerate(pts):
+                    cv2.circle(preview, tuple(pt[0]), 4, (0, 0, 255), -1)
+                    cv2.putText(
+                        preview,
+                        f"{i}.{j}",
+                        (pt[0][0] + 5, pt[0][1] - 5),
+                        cv2.FONT_HERSHEY_SIMPLEX,
+                        0.4,
+                        (255, 255, 0),
+                        1,
+                    )
+            cv2.imwrite(str(args.preview), preview)
+            print(f"Preview: {args.preview}")
+
+        print(f"Total polygons: {len(all_points)}")
+    else:
+        polygon = contours_to_polygon(contours, args.mode)
+        if polygon is None:
+            print("Error: No polygon generated", file=sys.stderr)
+            sys.exit(1)
+        padded = pad_polygon(polygon, args.padding)
+        simplified = simplify_to_target_points(padded, args.target_points)
+        clockwise = ensure_clockwise(simplified)
+
+        output_path = args.output
+        if output_path is None:
+            output_path = args.image.with_suffix(".tres")
+
+        uid = generate_godot_uid()
+        write_tres(clockwise, output_path, uid)
+
+        if args.preview:
+            write_preview(image, clockwise, args.preview)
+            print(f"Preview: {args.preview}")
+
+        print(f"Created: {output_path}")
+        print(f"Points: {len(clockwise.squeeze())}")
+        print(f"UID: uid://{uid}")
 
 
 if __name__ == "__main__":