|
| 1 | +# /// script |
| 2 | +# requires-python = ">=3.8" |
| 3 | +# dependencies = [ |
| 4 | +# "click", |
| 5 | +# "faker", |
| 6 | +# ] |
| 7 | +# /// |
| 8 | +import random |
| 9 | +import subprocess |
| 10 | +import json |
| 11 | +from pathlib import Path |
| 12 | +import click |
| 13 | +from faker import Faker |
| 14 | +import math |
| 15 | + |
| 16 | +fake = Faker() |
| 17 | + |
| 18 | +FRAME_WIDTH = 1280 |
| 19 | +FRAME_HEIGHT = 720 |
| 20 | + |
| 21 | +def create_random_video(file_path: Path, duration: int): |
| 22 | + """Create a random test video using ffmpeg (MP4 container, H.264 codec).""" |
| 23 | + cmd = [ |
| 24 | + "ffmpeg", "-y", |
| 25 | + "-f", "lavfi", "-i", "testsrc=size=1280x720:rate=30", |
| 26 | + "-t", str(duration), |
| 27 | + "-c:v", "libx264", |
| 28 | + "-pix_fmt", "yuv420p", |
| 29 | + str(file_path) |
| 30 | + ] |
| 31 | + subprocess.run(cmd, check=True) |
| 32 | + |
| 33 | +def extract_frames_from_video(video_path: Path, image_dir: Path): |
| 34 | + """Extract frames from a video and save as sequential JPG files.""" |
| 35 | + image_dir.mkdir(parents=True, exist_ok=True) |
| 36 | + cmd = [ |
| 37 | + "ffmpeg", "-y", |
| 38 | + "-i", str(video_path), |
| 39 | + str(image_dir / "frame_%04d.jpg") |
| 40 | + ] |
| 41 | + subprocess.run(cmd, check=True) |
| 42 | + |
| 43 | +def generate_star_points(cx, cy, r, spikes=5): |
| 44 | + """Generate points for a star polygon.""" |
| 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]) # close polygon |
| 53 | + return pts |
| 54 | + |
| 55 | +def generate_diamond_points(cx, cy, r): |
| 56 | + return [ |
| 57 | + [cx, cy - r], |
| 58 | + [cx + r, cy], |
| 59 | + [cx, cy + r], |
| 60 | + [cx - r, cy], |
| 61 | + [cx, cy - r], |
| 62 | + ] |
| 63 | + |
| 64 | +def generate_circle_points(cx, cy, r, segments=12): |
| 65 | + pts = [] |
| 66 | + for i in range(segments+1): |
| 67 | + angle = 2 * math.pi * i / segments |
| 68 | + x = cx + math.cos(angle) * r |
| 69 | + y = cy + math.sin(angle) * r |
| 70 | + pts.append([x, y]) |
| 71 | + return pts |
| 72 | + |
| 73 | +def generate_geometry(shape: str, cx: float, cy: float, size: float): |
| 74 | + """Return GeoJSON polygon of the shape centered at (cx, cy).""" |
| 75 | + if shape == "star": |
| 76 | + coords = generate_star_points(cx, cy, size) |
| 77 | + elif shape == "diamond": |
| 78 | + coords = generate_diamond_points(cx, cy, size) |
| 79 | + elif shape == "circle": |
| 80 | + coords = generate_circle_points(cx, cy, size) |
| 81 | + else: # rectangle fallback |
| 82 | + half = size |
| 83 | + coords = [ |
| 84 | + [cx-half, cy-half], |
| 85 | + [cx+half, cy-half], |
| 86 | + [cx+half, cy+half], |
| 87 | + [cx-half, cy+half], |
| 88 | + [cx-half, cy-half] |
| 89 | + ] |
| 90 | + return {'geojson': { |
| 91 | + "type": "FeatureCollection", |
| 92 | + "features": [ |
| 93 | + { |
| 94 | + "type" :"Feature", |
| 95 | + "geometry": { |
| 96 | + "type": "Polygon", |
| 97 | + "coordinates": [coords] |
| 98 | + }, |
| 99 | + "properties": { "key": "" } |
| 100 | + } |
| 101 | + ] |
| 102 | + }, |
| 103 | + 'coords': coords |
| 104 | + } |
| 105 | + |
| 106 | +def geometry_bounds(coords): |
| 107 | + """Calculate bounding box [x1, y1, x2, y2] from polygon points.""" |
| 108 | + xs = [pt[0] for pt in coords] |
| 109 | + ys = [pt[1] for pt in coords] |
| 110 | + return [min(xs), min(ys), max(xs), max(ys)] |
| 111 | + |
| 112 | +def generate_annotation_json(num_frames: int, output_file: Path): |
| 113 | + """Generate annotation JSON with moving/scaling geometry.""" |
| 114 | + num_tracks = random.randint(3, 5) |
| 115 | + tracks = {} |
| 116 | + |
| 117 | + for i in range(num_tracks): |
| 118 | + track_id = i |
| 119 | + shape_type = random.choice(["circle", "rectangle", "diamond", "star"]) |
| 120 | + begin = 0 |
| 121 | + end = num_frames - 1 |
| 122 | + |
| 123 | + # Initial position and motion |
| 124 | + x, y = random.randint(100, FRAME_WIDTH-100), random.randint(100, FRAME_HEIGHT-100) |
| 125 | + dx, dy = random.choice([-5, 5]), random.choice([-3, 3]) |
| 126 | + base_size = random.randint(40, 80) |
| 127 | + growth_rate = random.uniform(0.05, 0.15) |
| 128 | + |
| 129 | + features = [] |
| 130 | + for frame in range(num_frames): |
| 131 | + # Update position and bounce |
| 132 | + x += dx |
| 133 | + y += dy |
| 134 | + if x < 50 or x > FRAME_WIDTH-50: |
| 135 | + dx *= -1 |
| 136 | + x += dx |
| 137 | + if y < 50 or y > FRAME_HEIGHT-50: |
| 138 | + dy *= -1 |
| 139 | + y += dy |
| 140 | + |
| 141 | + # Smooth scaling |
| 142 | + scale = 0.5 * (1 + math.sin(growth_rate * frame)) |
| 143 | + size = base_size * (0.75 + 0.5 * scale) |
| 144 | + |
| 145 | + # Create moving geometry |
| 146 | + output_data = generate_geometry(shape_type, x, y, size) |
| 147 | + geom = output_data['geojson'] |
| 148 | + coords = output_data['coords'] |
| 149 | + bounds = geometry_bounds(coords) |
| 150 | + |
| 151 | + feature = { |
| 152 | + "frame": frame, |
| 153 | + "bounds": bounds, |
| 154 | + "keyframe": True, |
| 155 | + "geometry": geom |
| 156 | + } |
| 157 | + features.append(feature) |
| 158 | + |
| 159 | + tracks[str(track_id)] = { |
| 160 | + "id": track_id, |
| 161 | + "meta": {"shape": shape_type}, |
| 162 | + "attributes": {}, |
| 163 | + "confidencePairs": [[fake.word(), float(random.randrange(0, 100)/100)]], |
| 164 | + "begin": begin, |
| 165 | + "end": end, |
| 166 | + "features": features |
| 167 | + } |
| 168 | + |
| 169 | + annotation = { |
| 170 | + "tracks": tracks, |
| 171 | + "groups": {}, |
| 172 | + "version": 2 |
| 173 | + } |
| 174 | + with open(output_file, "w") as f: |
| 175 | + json.dump(annotation, f, indent=2) |
| 176 | + |
| 177 | +def create_video_content(base_dir: Path, max_videos: int, counter: dict, total: int): |
| 178 | + """Create videos and associated JSON annotations.""" |
| 179 | + if counter['count'] >= total: |
| 180 | + return |
| 181 | + num_videos = random.randint(1, max_videos) |
| 182 | + for _ in range(num_videos): |
| 183 | + if counter['count'] >= total: |
| 184 | + break |
| 185 | + duration = random.randint(5, 30) |
| 186 | + name = fake.word() + ".mp4" |
| 187 | + video_path = base_dir / name |
| 188 | + create_random_video(video_path, duration) |
| 189 | + counter['count'] += 1 |
| 190 | + |
| 191 | + # Generate annotation JSON |
| 192 | + generate_annotation_json(duration * 30, video_path.with_suffix(".json")) |
| 193 | + |
| 194 | +def create_image_sequence_content(base_dir: Path, counter: dict, total: int): |
| 195 | + """Create image sequence from a temporary video and generate JSON annotations.""" |
| 196 | + if counter['count'] >= total: |
| 197 | + return |
| 198 | + duration = random.randint(5, 30) |
| 199 | + tmp_video = base_dir / (fake.word() + "_tmp.mp4") |
| 200 | + create_random_video(tmp_video, duration) |
| 201 | + seq_folder = base_dir / (tmp_video.stem.replace("_tmp", "") + "_frames") |
| 202 | + extract_frames_from_video(tmp_video, seq_folder) |
| 203 | + tmp_video.unlink() |
| 204 | + counter['count'] += 1 |
| 205 | + |
| 206 | + # Generate annotation JSON for image sequence folder |
| 207 | + annotation_file = seq_folder / (seq_folder.stem + ".json") |
| 208 | + generate_annotation_json(duration * 30, annotation_file) |
| 209 | + |
| 210 | +def create_folder_structure(base_dir: Path, depth: int, max_depth: int, |
| 211 | + max_videos: int, counter: dict, total: int): |
| 212 | + """Recursively create folders with either videos or image sequences.""" |
| 213 | + if counter['count'] >= total: |
| 214 | + return |
| 215 | + |
| 216 | + content_type = random.choice(["video", "images"]) |
| 217 | + if content_type == "video": |
| 218 | + create_video_content(base_dir, max_videos, counter, total) |
| 219 | + else: |
| 220 | + create_image_sequence_content(base_dir, counter, total) |
| 221 | + |
| 222 | + if counter['count'] >= total: |
| 223 | + return |
| 224 | + |
| 225 | + num_subfolders = random.randint(0, 3) |
| 226 | + for _ in range(num_subfolders): |
| 227 | + if counter['count'] >= total: |
| 228 | + break |
| 229 | + subfolder = base_dir / fake.word() |
| 230 | + subfolder.mkdir(parents=True, exist_ok=True) |
| 231 | + if depth < max_depth: |
| 232 | + create_folder_structure(subfolder, depth+1, max_depth, max_videos, counter, total) |
| 233 | + else: |
| 234 | + leaf_type = random.choice(["video", "images"]) |
| 235 | + if leaf_type == "video": |
| 236 | + create_video_content(subfolder, max_videos, counter, total) |
| 237 | + else: |
| 238 | + create_image_sequence_content(subfolder, counter, total) |
| 239 | + |
| 240 | +@click.command() |
| 241 | +@click.option('--output', '-o', default='./sample', show_default=True, |
| 242 | + type=click.Path(file_okay=False), help="Base output directory") |
| 243 | +@click.option('--folders', '-f', default=3, show_default=True, |
| 244 | + help="Number of top-level folders to create") |
| 245 | +@click.option('--max-depth', '-d', default=2, show_default=True, |
| 246 | + help="Maximum subfolder depth") |
| 247 | +@click.option('--videos', '-v', default=2, show_default=True, |
| 248 | + help="Maximum videos per folder") |
| 249 | +@click.option('--total', '-t', default=10, show_default=True, |
| 250 | + help="Total number of datasets (videos or image sequences)") |
| 251 | +def main(output, folders, max_depth, videos, total): |
| 252 | + base_path = Path(output) |
| 253 | + base_path.mkdir(parents=True, exist_ok=True) |
| 254 | + |
| 255 | + counter = {'count': 0} |
| 256 | + click.echo(f"Generating up to {total} datasets in {base_path}...") |
| 257 | + for _ in range(folders): |
| 258 | + if counter['count'] >= total: |
| 259 | + break |
| 260 | + folder_path = base_path / fake.word() |
| 261 | + folder_path.mkdir(parents=True, exist_ok=True) |
| 262 | + create_folder_structure(folder_path, 1, max_depth, videos, counter, total) |
| 263 | + |
| 264 | + click.echo(f"Done! Created {counter['count']} datasets.") |
| 265 | + |
| 266 | +if __name__ == '__main__': |
| 267 | + main() |
0 commit comments