Files
ai-game-2/tools/draw_polygon.py
2026-03-09 09:22:21 -07:00

213 lines
6.1 KiB
Python

#!/usr/bin/env python3
"""Draw a polygon on an image using Pillow."""
import argparse
import sys
from pathlib import Path
try:
from PIL import Image, ImageDraw, ImageFilter
except ImportError:
print("Error: Pillow is required. Install with: pip install Pillow")
sys.exit(1)
COLOR_PALETTE = {
"red": "#FF0000",
"green": "#00FF00",
"blue": "#0000FF",
"yellow": "#FFFF00",
"cyan": "#00FFFF",
"magenta": "#FF00FF",
"white": "#FFFFFF",
"black": "#000000",
"orange": "#FFA500",
"purple": "#800080",
"brown": "#A52A2A",
"pink": "#FFC0CB",
"gray": "#808080",
"grey": "#808080",
"lime": "#00FF00",
"navy": "#000080",
"teal": "#008080",
"maroon": "#800000",
}
def parse_color(color_str: str) -> str:
"""Parse color string to hex."""
color_lower = color_str.lower()
if color_lower in COLOR_PALETTE:
return COLOR_PALETTE[color_lower]
if color_str.startswith("#") and len(color_str) in [7, 9]:
return color_str
raise ValueError(
f"Invalid color '{color_str}'. Use a name (e.g., red, blue) or hex code (e.g., #FF0000)."
)
def parse_points(points_str: str) -> list[tuple[float, float]]:
"""Parse points string into list of (x, y) tuples."""
points = []
for point in points_str.split():
if "," not in point:
raise ValueError(f"Invalid point format: '{point}'. Use 'x,y' format.")
try:
x, y = point.split(",")
points.append((float(x.strip()), float(y.strip())))
except ValueError:
raise ValueError(f"Invalid point format: '{point}'. Use 'x,y' format.")
return points
def convert_to_pixels(
points: list[tuple[float, float]], width: int, height: int, absolute: bool
) -> list[tuple[int, int]]:
"""Convert points to pixel coordinates."""
if absolute:
return [(int(x), int(y)) for x, y in points]
else:
return [
(int(x * width), int(y * height))
for x, y in points
if 0 <= x <= 1 and 0 <= y <= 1
]
def draw_polygon_on_image(
image_path: Path,
points: list[tuple[float, float]],
color: str,
thickness: int,
fill: bool,
absolute: bool,
save_path: Path | None,
) -> None:
"""Draw polygon on image."""
image = Image.open(image_path).convert("RGBA")
width, height = image.size
pixel_points = convert_to_pixels(points, width, height, absolute)
if len(pixel_points) < 3:
raise ValueError("Need at least 3 points to draw a polygon.")
draw = ImageDraw.Draw(image)
hex_color = parse_color(color)
draw.polygon(pixel_points, fill=hex_color if fill else None, outline=hex_color, width=thickness)
if fill:
alpha = 128
fill_color = hex_color[:7] + f"{alpha:02X}"
draw.polygon(pixel_points, fill=fill_color, outline=hex_color, width=thickness)
if save_path:
image.save(save_path)
print(f"Saved: {save_path}")
else:
image.show()
def main():
parser = argparse.ArgumentParser(
description="Draw a polygon on an image using Pillow.",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Examples:
# Draw a triangle using percentages (default)
python draw_polygon.py scene.png "0.5,0.1 0.9,0.9 0.1,0.9"
# Draw with pixel coordinates (absolute mode)
python draw_polygon.py scene.png "100,50 400,50 250,300" --absolute
# Save to file with blue color
python draw_polygon.py scene.png "0.2,0.2 0.8,0.2 0.8,0.8 0.2,0.8" --color blue --save output.png
# Fill polygon semi-transparently
python draw_polygon.py scene.png "0.3,0.3 0.7,0.3 0.7,0.7 0.3,0.7" --fill --color red
# Draw with custom hex color and thickness
python draw_polygon.py scene.png "0.1,0.1 0.9,0.1 0.9,0.9 0.1,0.9" --color #00FF00 --thickness 5
Coordinate Formats:
Percentage (default): 0.0 to 1.0, where 0.0,0.0 is top-left and 1.0,1.0 is bottom-right
Absolute (with --absolute): Actual pixel coordinates from 0 to image width/height
Color Formats:
Named colors: red, green, blue, yellow, cyan, magenta, white, black, orange, purple, brown, pink, gray, grey, lime, navy, teal, maroon
Hex codes: #FF0000, #00FF00FF, etc.
""",
)
parser.add_argument(
"image",
type=Path,
help="Path to input image file",
)
parser.add_argument(
"points",
type=str,
help="Space-separated list of x,y coordinates (e.g., '0.1,0.1 0.9,0.1 0.5,0.9'). Use percentages (0.0-1.0) by default, or --absolute for pixel coordinates.",
)
parser.add_argument(
"--color",
type=str,
default="red",
help="Polygon color (default: red). Use named colors (red, blue, green, etc.) or hex codes (#FF0000).",
)
parser.add_argument(
"--absolute",
action="store_true",
help="Use pixel coordinates instead of percentages",
)
parser.add_argument(
"--save",
type=Path,
default=None,
help="Save output to file path instead of displaying",
)
parser.add_argument(
"--thickness",
type=int,
default=2,
help="Line thickness in pixels (default: 2)",
)
parser.add_argument(
"--fill",
action="store_true",
help="Fill polygon with semi-transparent color",
)
args = parser.parse_args()
if not args.image.exists():
print(f"Error: Image not found: {args.image}", file=sys.stderr)
sys.exit(1)
try:
points = parse_points(args.points)
if len(points) < 3:
print(f"Error: Need at least 3 points. Got {len(points)}.", file=sys.stderr)
sys.exit(1)
draw_polygon_on_image(
image_path=args.image,
points=points,
color=args.color,
thickness=args.thickness,
fill=args.fill,
absolute=args.absolute,
save_path=args.save,
)
except ValueError as e:
print(f"Error: {e}", file=sys.stderr)
sys.exit(1)
except Exception as e:
print(f"Error: {e}", file=sys.stderr)
sys.exit(1)
if __name__ == "__main__":
main()