57 lines
2.1 KiB
Python
57 lines
2.1 KiB
Python
from __future__ import annotations
|
|
|
|
import math
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from PIL import Image, ImageDraw
|
|
|
|
RGB = {
|
|
"red": (220, 55, 55),
|
|
"blue": (55, 100, 220),
|
|
"green": (45, 165, 85),
|
|
"yellow": (235, 195, 45),
|
|
}
|
|
|
|
|
|
def _polygon(cx: int, cy: int, radius: int, sides: int, rotation: float = -math.pi / 2):
|
|
return [
|
|
(
|
|
cx + radius * math.cos(rotation + (2 * math.pi * index / sides)),
|
|
cy + radius * math.sin(rotation + (2 * math.pi * index / sides)),
|
|
)
|
|
for index in range(sides)
|
|
]
|
|
|
|
|
|
def render_scene(record: dict[str, Any], output_path: Path, image_size: int = 224) -> None:
|
|
image = Image.new("RGB", (image_size, image_size), (248, 248, 248))
|
|
draw = ImageDraw.Draw(image)
|
|
scale = image_size / 224
|
|
radius = int(28 * scale)
|
|
for item in record["objects"]:
|
|
x, y = int(item["x"] * scale), int(item["y"] * scale)
|
|
color = RGB[item["color"]]
|
|
box = (x - radius, y - radius, x + radius, y + radius)
|
|
if item["shape"] == "circle":
|
|
draw.ellipse(box, fill=color, outline=(20, 20, 20), width=max(1, int(2 * scale)))
|
|
elif item["shape"] == "square":
|
|
draw.rectangle(box, fill=color, outline=(20, 20, 20), width=max(1, int(2 * scale)))
|
|
elif item["shape"] == "triangle":
|
|
draw.polygon(_polygon(x, y, radius, 3), fill=color, outline=(20, 20, 20))
|
|
else:
|
|
points = []
|
|
for index in range(10):
|
|
current_radius = radius if index % 2 == 0 else radius * 0.45
|
|
angle = -math.pi / 2 + index * math.pi / 5
|
|
points.append(
|
|
(
|
|
x + current_radius * math.cos(angle),
|
|
y + current_radius * math.sin(angle),
|
|
)
|
|
)
|
|
draw.polygon(points, fill=color, outline=(20, 20, 20))
|
|
draw.text((x - radius, y + radius + int(4 * scale)), item["object_id"], fill=(20, 20, 20))
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
image.save(output_path)
|