|
| 1 | +# /// script |
| 2 | +# requires-python = ">=3.8" |
| 3 | +# dependencies = [ |
| 4 | +# "click", |
| 5 | +# "faker", |
| 6 | +# ] |
| 7 | +# /// |
| 8 | +import subprocess |
| 9 | +import shutil |
| 10 | +import random |
| 11 | +import json |
| 12 | +import math |
| 13 | +from pathlib import Path |
| 14 | +import click |
| 15 | +from faker import Faker |
| 16 | + |
| 17 | +fake = Faker() |
| 18 | + |
| 19 | +FRAME_WIDTH = 1280 |
| 20 | +FRAME_HEIGHT = 720 |
| 21 | +FPS = 30 |
| 22 | + |
| 23 | + |
| 24 | +def create_base_video(file_path: Path, duration: int): |
| 25 | + """Create a single base test video using ffmpeg (MP4 container, H.264 codec).""" |
| 26 | + cmd = [ |
| 27 | + "ffmpeg", "-y", |
| 28 | + "-f", "lavfi", "-i", f"testsrc=size={FRAME_WIDTH}x{FRAME_HEIGHT}:rate={FPS}", |
| 29 | + "-t", str(duration), |
| 30 | + "-c:v", "libx264", |
| 31 | + "-pix_fmt", "yuv420p", |
| 32 | + str(file_path) |
| 33 | + ] |
| 34 | + subprocess.run(cmd, check=True) |
| 35 | + |
| 36 | + |
| 37 | +def copy_video_to_new_location(base_video: Path, target_path: Path): |
| 38 | + """Copy the base video to a new file path.""" |
| 39 | + shutil.copy2(base_video, target_path) |
| 40 | + |
| 41 | + |
| 42 | +# ---- Geometry + Annotations ---- |
| 43 | + |
| 44 | +def generate_star_points(cx, cy, r, spikes=5): |
| 45 | + pts = [] |
| 46 | + angle = math.pi / spikes |
| 47 | + for i in range(2 * spikes): |
| 48 | + radius = r if i % 2 == 0 else r / 2 |
| 49 | + x = cx + math.cos(i * angle) * radius |
| 50 | + y = cy + math.sin(i * angle) * radius |
| 51 | + pts.append([x, y]) |
| 52 | + pts.append(pts[0]) |
| 53 | + return pts |
| 54 | + |
| 55 | + |
| 56 | +def generate_diamond_points(cx, cy, r): |
| 57 | + return [ |
| 58 | + [cx, cy - r], |
| 59 | + [cx + r, cy], |
| 60 | + [cx, cy + r], |
| 61 | + [cx - r, cy], |
| 62 | + [cx, cy - r], |
| 63 | + ] |
| 64 | + |
| 65 | + |
| 66 | +def generate_circle_points(cx, cy, r, segments=12): |
| 67 | + pts = [] |
| 68 | + for i in range(segments + 1): |
| 69 | + angle = 2 * math.pi * i / segments |
| 70 | + x = cx + math.cos(angle) * r |
| 71 | + y = cy + math.sin(angle) * r |
| 72 | + pts.append([x, y]) |
| 73 | + return pts |
| 74 | + |
| 75 | + |
| 76 | +def generate_geometry(shape: str, cx: float, cy: float, size: float): |
| 77 | + if shape == "star": |
| 78 | + coords = generate_star_points(cx, cy, size) |
| 79 | + elif shape == "diamond": |
| 80 | + coords = generate_diamond_points(cx, cy, size) |
| 81 | + elif shape == "circle": |
| 82 | + coords = generate_circle_points(cx, cy, size) |
| 83 | + else: # rectangle fallback |
| 84 | + half = size |
| 85 | + coords = [ |
| 86 | + [cx-half, cy-half], |
| 87 | + [cx+half, cy-half], |
| 88 | + [cx+half, cy+half], |
| 89 | + [cx-half, cy+half], |
| 90 | + [cx-half, cy-half] |
| 91 | + ] |
| 92 | + return {'geojson': { |
| 93 | + "type": "FeatureCollection", |
| 94 | + "features": [ |
| 95 | + { |
| 96 | + "type": "Feature", |
| 97 | + "geometry": { |
| 98 | + "type": "Polygon", |
| 99 | + "coordinates": [coords] |
| 100 | + }, |
| 101 | + "properties": { "key": "" } |
| 102 | + } |
| 103 | + ] |
| 104 | + }, |
| 105 | + 'coords': coords |
| 106 | + } |
| 107 | + |
| 108 | + |
| 109 | +def geometry_bounds(coords): |
| 110 | + xs = [pt[0] for pt in coords] |
| 111 | + ys = [pt[1] for pt in coords] |
| 112 | + return [min(xs), min(ys), max(xs), max(ys)] |
| 113 | + |
| 114 | + |
| 115 | +def generate_annotation_json(num_frames: int, output_file: Path): |
| 116 | + """Generate annotation JSON with moving/scaling geometry.""" |
| 117 | + num_tracks = random.randint(3, 5) |
| 118 | + tracks = {} |
| 119 | + |
| 120 | + for i in range(num_tracks): |
| 121 | + track_id = i |
| 122 | + shape_type = random.choice(["circle", "rectangle", "diamond", "star"]) |
| 123 | + begin = 0 |
| 124 | + end = num_frames - 1 |
| 125 | + |
| 126 | + x, y = random.randint(100, FRAME_WIDTH-100), random.randint(100, FRAME_HEIGHT-100) |
| 127 | + dx, dy = random.choice([-5, 5]), random.choice([-3, 3]) |
| 128 | + base_size = random.randint(40, 80) |
| 129 | + growth_rate = random.uniform(0.05, 0.15) |
| 130 | + |
| 131 | + features = [] |
| 132 | + for frame in range(num_frames): |
| 133 | + x += dx |
| 134 | + y += dy |
| 135 | + if x < 50 or x > FRAME_WIDTH-50: |
| 136 | + dx *= -1 |
| 137 | + x += dx |
| 138 | + if y < 50 or y > FRAME_HEIGHT-50: |
| 139 | + dy *= -1 |
| 140 | + y += dy |
| 141 | + |
| 142 | + scale = 0.5 * (1 + math.sin(growth_rate * frame)) |
| 143 | + size = base_size * (0.75 + 0.5 * scale) |
| 144 | + |
| 145 | + output_data = generate_geometry(shape_type, x, y, size) |
| 146 | + geom = output_data['geojson'] |
| 147 | + coords = output_data['coords'] |
| 148 | + bounds = geometry_bounds(coords) |
| 149 | + |
| 150 | + feature = { |
| 151 | + "frame": frame, |
| 152 | + "bounds": bounds, |
| 153 | + "keyframe": True, |
| 154 | + "geometry": geom |
| 155 | + } |
| 156 | + features.append(feature) |
| 157 | + |
| 158 | + tracks[str(track_id)] = { |
| 159 | + "id": track_id, |
| 160 | + "meta": {"shape": shape_type}, |
| 161 | + "attributes": {}, |
| 162 | + "confidencePairs": [[fake.word(), float(random.randrange(0, 100)/100)]], |
| 163 | + "begin": begin, |
| 164 | + "end": end, |
| 165 | + "features": features |
| 166 | + } |
| 167 | + |
| 168 | + annotation = { |
| 169 | + "tracks": tracks, |
| 170 | + "groups": {}, |
| 171 | + "version": 2 |
| 172 | + } |
| 173 | + with open(output_file, "w") as f: |
| 174 | + json.dump(annotation, f, indent=2) |
| 175 | + |
| 176 | + |
| 177 | +# ---- Main orchestration ---- |
| 178 | + |
| 179 | +@click.command() |
| 180 | +@click.option('--output', '-o', default='./sample', show_default=True, |
| 181 | + type=click.Path(file_okay=False), help="Base output directory") |
| 182 | +@click.option('--folders', '-f', default=3, show_default=True, |
| 183 | + help="Number of top-level folders to create") |
| 184 | +@click.option('--max-depth', '-d', default=2, show_default=True, |
| 185 | + help="Maximum subfolder depth") |
| 186 | +@click.option('--videos', '-v', default=4, show_default=True, |
| 187 | + help="Number of videos per folder") |
| 188 | +@click.option('--total', '-t', default=50, show_default=True, |
| 189 | + help="Total number of videos to create") |
| 190 | +@click.option('--duration', default=10, show_default=True, |
| 191 | + help="Duration (in seconds) of each video") |
| 192 | +def main(output, folders, max_depth, videos, total, duration): |
| 193 | + base_path = Path(output) |
| 194 | + base_path.mkdir(parents=True, exist_ok=True) |
| 195 | + |
| 196 | + # Create one base video |
| 197 | + base_video = base_path / "base_video.mp4" |
| 198 | + click.echo(f"Creating base video of {duration}s...") |
| 199 | + create_base_video(base_video, duration) |
| 200 | + |
| 201 | + num_frames = duration * FPS |
| 202 | + count = 0 |
| 203 | + for f in range(folders): |
| 204 | + if count >= total: |
| 205 | + break |
| 206 | + folder_path = base_path / fake.word() |
| 207 | + folder_path.mkdir(parents=True, exist_ok=True) |
| 208 | + |
| 209 | + stack = [(folder_path, 1)] |
| 210 | + while stack and count < total: |
| 211 | + current_path, depth = stack.pop() |
| 212 | + |
| 213 | + for _ in range(videos): |
| 214 | + if count >= total: |
| 215 | + break |
| 216 | + name = fake.word() |
| 217 | + video_path = current_path / f"{name}.mp4" |
| 218 | + json_path = current_path / f"{name}.json" |
| 219 | + |
| 220 | + # Copy base video and create new annotations |
| 221 | + copy_video_to_new_location(base_video, video_path) |
| 222 | + generate_annotation_json(num_frames, json_path) |
| 223 | + count += 1 |
| 224 | + |
| 225 | + if depth < max_depth: |
| 226 | + for _ in range(2): # up to 2 subfolders each |
| 227 | + subfolder = current_path / fake.word() |
| 228 | + subfolder.mkdir(parents=True, exist_ok=True) |
| 229 | + stack.append((subfolder, depth+1)) |
| 230 | + |
| 231 | + click.echo(f"Done! Created {count} videos (all {duration}s long, with annotations).") |
| 232 | + |
| 233 | + |
| 234 | +if __name__ == '__main__': |
| 235 | + main() |
0 commit comments