Files
ai-game-2/tools/mask_to_polygon.py

285 lines
7.5 KiB
Python
Executable File

#!/usr/bin/env python3
"""
Convert a black/white mask image to a Godot .tres polygon resource.
"""
import argparse
import secrets
import sys
from pathlib import Path
import cv2
import numpy as np
from shapely.geometry import Polygon
from shapely.ops import unary_union
ALPHABET = "0123456789abcdefghijklmnopqrstuvwxyz"
def to_base36(num: int) -> str:
"""Convert integer to lowercase base36 (like Godot)."""
if num == 0:
return "0"
chars = []
while num:
num, rem = divmod(num, 36)
chars.append(ALPHABET[rem])
return "".join(reversed(chars))
def generate_godot_uid() -> str:
"""Generate a random Godot-compatible UID."""
value = secrets.randbits(64)
return to_base36(value)
def find_largest_contour(contours: list) -> np.ndarray:
"""Find the largest contour by area."""
if not contours:
return None
return max(contours, key=cv2.contourArea)
def contours_to_polygon(contours: list, mode: str) -> np.ndarray:
"""
Convert contours to a single polygon.
Args:
contours: List of contours from cv2.findContours
mode: "convex_hull" or "largest_only"
Returns:
Single contour as numpy array
"""
if not contours:
return None
if mode == "largest_only":
return find_largest_contour(contours)
if mode == "convex_hull":
all_points = np.vstack(contours)
return cv2.convexHull(all_points)
raise ValueError(f"Unknown mode: {mode}")
def simplify_to_target_points(contour: np.ndarray, target_points: int) -> np.ndarray:
"""
Simplify a contour to approximately target number of points.
Uses Douglas-Peucker algorithm with iterative epsilon adjustment.
"""
if len(contour) <= target_points:
return contour
contour = contour.squeeze()
if len(contour.shape) == 1:
return contour.reshape(-1, 1, 2)
perimeter = cv2.arcLength(contour, True)
epsilon_low = 0.0
epsilon_high = perimeter * 0.1
best_result = contour
for _ in range(50):
epsilon_mid = (epsilon_low + epsilon_high) / 2
simplified = cv2.approxPolyDP(contour, epsilon_mid, True)
if len(simplified) == target_points:
return simplified
if len(simplified) > target_points:
epsilon_low = epsilon_mid
else:
epsilon_high = epsilon_mid
best_result = simplified
if abs(len(simplified) - target_points) <= 1:
best_result = simplified
break
return best_result
def pad_polygon(points: np.ndarray, padding: float) -> np.ndarray:
"""
Expand polygon outward by padding pixels using Shapely buffer.
"""
if padding <= 0:
return points
points_2d = points.squeeze()
if len(points_2d.shape) == 1:
return points
polygon = Polygon(points_2d)
if not polygon.is_valid:
polygon = polygon.buffer(0)
buffered = polygon.buffer(padding)
if buffered.is_empty:
return points
if buffered.geom_type == "MultiPolygon":
buffered = max(buffered.geoms, key=lambda p: p.area)
exterior_coords = list(buffered.exterior.coords)[:-1]
return np.array(exterior_coords, dtype=np.float32).reshape(-1, 1, 2)
def ensure_clockwise(points: np.ndarray) -> np.ndarray:
"""
Ensure points are in clockwise order (visually).
In image coordinates (Y-down), clockwise means positive area.
"""
points_2d = points.squeeze()
if len(points_2d.shape) == 1:
return points
area = 0.0
n = len(points_2d)
for i in range(n):
j = (i + 1) % n
area += points_2d[i][0] * points_2d[j][1]
area -= points_2d[j][0] * points_2d[i][1]
if area < 0:
return points[::-1]
return points
def write_preview(
original_image: np.ndarray, points: np.ndarray, output_path: Path
) -> None:
"""
Write a preview image showing the polygon drawn on the mask.
"""
preview = cv2.cvtColor(original_image, cv2.COLOR_GRAY2BGR)
pts = points.squeeze().astype(np.int32).reshape(-1, 1, 2)
cv2.polylines(preview, [pts], True, (0, 255, 0), 2)
for i, pt in enumerate(pts):
cv2.circle(preview, tuple(pt[0]), 4, (0, 0, 255), -1)
cv2.putText(
preview,
str(i),
(pt[0][0] + 5, pt[0][1] - 5),
cv2.FONT_HERSHEY_SIMPLEX,
0.4,
(255, 255, 0),
1,
)
cv2.imwrite(str(output_path), preview)
def write_tres(points: np.ndarray, output_path: Path, uid: str) -> None:
"""
Write polygon points to a Godot .tres resource file.
"""
points_2d = points.squeeze()
if len(points_2d.shape) == 1:
raise ValueError("Invalid points array")
coords = []
for pt in points_2d:
coords.extend([round(pt[0], 2), round(pt[1], 2)])
coords_str = ", ".join(str(c) for c in coords)
script_uid = "uid://dtemboas3bi8y"
content = f"""[gd_resource type="Resource" script_class="PolygonPointsResource" format=3 uid="uid://{uid}"]
[ext_resource type="Script" uid="{script_uid}" path="res://PolygonPointsResource.gd" id="1_ppr"]
[resource]
script = ExtResource("1_ppr")
points = PackedVector2Array({coords_str})
metadata/_custom_type_script = "{script_uid}"
"""
output_path.write_text(content)
def main():
parser = argparse.ArgumentParser(
description="Convert a black/white mask to a Godot polygon .tres resource"
)
parser.add_argument("image", type=Path, help="Input mask image (black/white)")
parser.add_argument(
"--padding",
type=float,
default=0,
help="Pixels to expand polygon outward (default: 0)",
)
parser.add_argument(
"--target-points",
type=int,
default=10,
help="Target number of polygon points (default: 10)",
)
parser.add_argument(
"--output",
type=Path,
default=None,
help="Output .tres file path (default: <image_stem>.tres)",
)
parser.add_argument(
"--mode",
choices=["convex_hull", "largest_only"],
default="convex_hull",
help="How to handle multiple regions (default: convex_hull)",
)
parser.add_argument(
"--preview",
type=Path,
default=None,
help="Output preview PNG showing polygon on mask",
)
args = parser.parse_args()
if not args.image.exists():
print(f"Error: Image not found: {args.image}", file=sys.stderr)
sys.exit(1)
image = cv2.imread(str(args.image), cv2.IMREAD_GRAYSCALE)
if image is None:
print(f"Error: Could not load image: {args.image}", file=sys.stderr)
sys.exit(1)
_, binary = cv2.threshold(image, 127, 255, cv2.THRESH_BINARY)
contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
if not contours:
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)
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__":
main()