Spaces:
Build error
Build error
File size: 3,202 Bytes
4fa8bcb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 |
from PIL import Image, ImageDraw, ImageFont
import os
from utils.normalize_interpolate import normalize_interpolate
from utils.path_utils import resolve_path
def generate_image(prompt, frame_number, normalized_pos, output_folder="frames", ramp_start=(64, 64),ramp_end=(448, 448), img_size=(512, 512), ball_radius=15):
"""
Args:
prompt (str): Text prompt (optional, for annotation)
frame_number (int): Frame index
ball_position_normalized (float): 0.0 (top) to 1.0 (bottom) along the ramp
output_folder (str): Where to save images
ramp_start: start of ramp in pixel coords
ramp_end: end of ramp in pixel coords
"""
os.makedirs(output_folder, exist_ok=True)
img_size = (512, 512)
img = Image.new('RGB', img_size, color=(255, 255, 255))
d = ImageDraw.Draw(img)
# Position of the ball
x, y = normalize_interpolate(ramp_start, ramp_end, normalized_pos)
# Draw the ramp line
d.line([ramp_start, ramp_end], fill=(150, 150, 150), width=3)
# Draw ball
d.ellipse(
[(x - ball_radius, y - ball_radius), (x + ball_radius, y + ball_radius)],
fill=(255, 0, 0), outline=(0, 0, 0)
)
try:
font = ImageFont.truetype("arial.ttf", 20)
except:
font = ImageFont.load_default()
d.text((10, 10), f"Frame {frame_number}: {prompt[:30]}", fill=(0, 0, 0), font=font)
file_path = resolve_path(f"frame_{frame_number:03}.png", subdir=output_folder, write_mode=True)
img.save(file_path)
return file_path
def draw_scene(
entities,
frame_number,
prompt="",
output_folder="frames",
img_size=(512, 512),
background_color=(255, 255, 255)
):
"""
Draws a general physics frame from a list of entities.
Args:
entities (List[Dict]): Each entity has a type, position, shape, etc.
frame_number (int): Index of the frame
prompt (str): Annotation label
output_folder (str): Directory to save frames
img_size (Tuple[int, int]): Image size
background_color (Tuple[int, int, int]): RGB color
"""
img = Image.new('RGB', img_size, color=background_color)
draw = ImageDraw.Draw(img)
for entity in entities:
etype = entity.get("type")
pos = entity.get("position") # (x, y) in pixels
radius = entity.get("radius", 5)
color = entity.get("color", (255, 0, 0))
if etype == "particle":
x, y = pos
draw.ellipse(
[(x - radius, y - radius), (x + radius, y + radius)],
fill=color, outline=(0, 0, 0)
)
elif etype == "line":
draw.line(entity["points"], fill=color, width=entity.get("width", 2))
elif etype == "text":
try:
font = ImageFont.truetype("arial.ttf", entity.get("size", 16))
except:
font = ImageFont.load_default()
draw.text(pos, entity["text"], fill=color, font=font)
# Future types: vector fields, heatmaps, force arrows, etc.
path = resolve_path(f"frame_{frame_number:03}.png", subdir=output_folder, write_mode=True)
img.save(path)
return path |