Setting the file. One moment. 16.11
Validate Atlas
def component_bounds
— line 251
This file
- Number
- 16.6
- Position
- 6 of 13
- Type
- Python
- Size
- 14 KB
- Lines
- 408
scripts/extract_strip_frames.py
Python·408 lines·14 KB
14CELL_WIDTH = 192
15CELL_HEIGHT = 208
16ROW_FRAME_COUNTS = {
17 "idle": 6,
18 "running-right": 8,
19 "running-left": 8,
20 "waving": 4,
21 "jumping": 5,
22 "failed": 8,
23 "waiting": 6,
24 "running": 6,
25 "review": 6,
26}
27
28
29def parse_states(raw: str) -> list[str]:
30 if raw.strip().lower() == "all":
31 return list(ROW_FRAME_COUNTS)
32 states = [item.strip() for item in raw.split(",") if item.strip()]
33 unknown = sorted(set(states) - set(ROW_FRAME_COUNTS))
34 if unknown:
35 raise SystemExit(f"unknown state(s): {', '.join(unknown)}")
36 return states
37
38
39def parse_hex_color(value: str) -> tuple[int, int, int]:
40 if not re.fullmatch(r"#[0-9a-fA-F]{6}", value):
41 raise SystemExit(f"invalid chroma key color: {value}; expected #RRGGBB")
42 return tuple(int(value[index : index + 2], 16) for index in (1, 3, 5))
43
44
45def load_chroma_key(decoded_dir: Path, override: str | None) -> tuple[int, int, int]:
46 if override:
47 return parse_hex_color(override)
48 request_path = decoded_dir.parent / "pet_request.json"
49 if request_path.is_file():
50 request = json.loads(request_path.read_text(encoding="utf-8"))
51 chroma_key = request.get("chroma_key")
52 if isinstance(chroma_key, dict) and isinstance(chroma_key.get("hex"), str):
53 return parse_hex_color(chroma_key["hex"])
54 return parse_hex_color("#00FF00")
55
56
57def color_distance(
58 red: int,
59 green: int,
60 blue: int,
61 key: tuple[int, int, int],
62) -> float:
63 return math.sqrt((red - key[0]) ** 2 + (green - key[1]) ** 2 + (blue - key[2]) ** 2)
64
65
66def remove_chroma_background(
67 image: Image.Image,
68 chroma_key: tuple[int, int, int],
69 threshold: float,
70) -> Image.Image:
71 rgba = image.convert("RGBA")
72 pixels = rgba.load()
73 for y in range(rgba.height):
74 for x in range(rgba.width):
75 red, green, blue, alpha = pixels[x, y]
76 if color_distance(red, green, blue, chroma_key) <= threshold:
77 pixels[x, y] = (0, 0, 0, 0)
78 return rgba
79
80
81def fit_to_cell(image: Image.Image) -> Image.Image:
82 bbox = image.getbbox()
83 target = Image.new("RGBA", (CELL_WIDTH, CELL_HEIGHT), (0, 0, 0, 0))
84 if bbox is None:
85 return target
86
87 sprite = image.crop(bbox)
88 max_width = CELL_WIDTH - 10
89 max_height = CELL_HEIGHT - 10
90 scale = min(max_width / sprite.width, max_height / sprite.height, 1.0)
91 if scale != 1.0:
92 sprite = sprite.resize(
93 (max(1, round(sprite.width * scale)), max(1, round(sprite.height * scale))),
94 Image.Resampling.LANCZOS,
95 )
96 left = (CELL_WIDTH - sprite.width) // 2
97 top = (CELL_HEIGHT - sprite.height) // 2
98 target.alpha_composite(sprite, (left, top))
99 return target
100
101
102def fit_viewport_to_cell(image: Image.Image) -> Image.Image:
103 target = Image.new("RGBA", (CELL_WIDTH, CELL_HEIGHT), (0, 0, 0, 0))
104 if image.getbbox() is None:
105 return target
106
107 viewport = image.copy()
108 max_width = CELL_WIDTH - 10
109 max_height = CELL_HEIGHT - 10
110 scale = min(max_width / viewport.width, max_height / viewport.height, 1.0)
111 if scale != 1.0:
112 viewport = viewport.resize(
113 (max(1, round(viewport.width * scale)), max(1, round(viewport.height * scale))),
114 Image.Resampling.LANCZOS,
115 )
116 left = (CELL_WIDTH - viewport.width) // 2
117 top = (CELL_HEIGHT - viewport.height) // 2
118 target.alpha_composite(viewport, (left, top))
119 return target
120
121
122def connected_components(image: Image.Image) -> list[dict[str, object]]:
123 alpha = image.getchannel("A")
124 width, height = image.size
125 data = alpha.tobytes()
126 visited = bytearray(width * height)
127 components: list[dict[str, object]] = []
128
129 for start, alpha_value in enumerate(data):
130 if alpha_value <= 16 or visited[start]:
131 continue
132
133 stack = [start]
134 visited[start] = 1
135 pixels: list[int] = []
136 min_x = width
137 min_y = height
138 max_x = 0
139 max_y = 0
140
141 while stack:
142 current = stack.pop()
143 pixels.append(current)
144 x = current % width
145 y = current // width
146 min_x = min(min_x, x)
147 min_y = min(min_y, y)
148 max_x = max(max_x, x)
149 max_y = max(max_y, y)
150
151 if x > 0:
152 neighbor = current - 1
153 if not visited[neighbor] and data[neighbor] > 16:
154 visited[neighbor] = 1
155 stack.append(neighbor)
156 if x + 1 < width:
157 neighbor = current + 1
158 if not visited[neighbor] and data[neighbor] > 16:
159 visited[neighbor] = 1
160 stack.append(neighbor)
161 if y > 0:
162 neighbor = current - width
163 if not visited[neighbor] and data[neighbor] > 16:
164 visited[neighbor] = 1
165 stack.append(neighbor)
166 if y + 1 < height:
167 neighbor = current + width
168 if not visited[neighbor] and data[neighbor] > 16:
169 visited[neighbor] = 1
170 stack.append(neighbor)
171
172 components.append(
173 {
174 "pixels": pixels,
175 "area": len(pixels),
176 "bbox": (min_x, min_y, max_x + 1, max_y + 1),
177 "center_x": (min_x + max_x + 1) / 2,
178 }
179 )
180
181 return components
182
183
184def component_group_image(
185 source: Image.Image,
186 components: list[dict[str, object]],
187 padding: int = 4,
188) -> Image.Image:
189 width, height = source.size
190 min_x = max(0, min(component["bbox"][0] for component in components) - padding)
191 min_y = max(0, min(component["bbox"][1] for component in components) - padding)
192 max_x = min(width, max(component["bbox"][2] for component in components) + padding)
193 max_y = min(height, max(component["bbox"][3] for component in components) + padding)
194
195 output = Image.new("RGBA", (max_x - min_x, max_y - min_y), (0, 0, 0, 0))
196 source_pixels = source.load()
197 output_pixels = output.load()
198 for component in components:
199 for pixel_index in component["pixels"]:
200 x = pixel_index % width
201 y = pixel_index // width
202 output_pixels[x - min_x, y - min_y] = source_pixels[x, y]
203 return output
204
205
206def component_frame_groups(
207 strip: Image.Image,
208 frame_count: int,
209) -> list[list[dict[str, object]]] | None:
210 components = connected_components(strip)
211 if not components:
212 return None
213
214 largest_area = max(component["area"] for component in components)
215 seed_threshold = max(120, largest_area * 0.20)
216 seeds = [component for component in components if component["area"] >= seed_threshold]
217 if len(seeds) < frame_count:
218 seeds = sorted(components, key=lambda component: component["area"], reverse=True)[
219 :frame_count
220 ]
221 if len(seeds) < frame_count:
222 return None
223
224 seeds = sorted(
225 sorted(seeds, key=lambda component: component["area"], reverse=True)[:frame_count],
226 key=lambda component: component["center_x"],
227 )
228 seed_ids = {id(seed) for seed in seeds}
229 groups: list[list[dict[str, object]]] = [[seed] for seed in seeds]
230 noise_threshold = max(12, largest_area * 0.002)
231
232 for component in components:
233 if id(component) in seed_ids or component["area"] < noise_threshold:
234 continue
235 nearest_index = min(
236 range(len(seeds)),
237 key=lambda index: abs(seeds[index]["center_x"] - component["center_x"]),
238 )
239 groups[nearest_index].append(component)
240
241 return groups
242
243
244def extract_component_frames(strip: Image.Image, frame_count: int) -> list[Image.Image] | None:
245 groups = component_frame_groups(strip, frame_count)
246 if groups is None:
247 return None
248 return [fit_to_cell(component_group_image(strip, group)) for group in groups]
249
250
251def component_bounds(components: list[dict[str, object]]) -> tuple[int, int, int, int]:
252 return (
253 min(component["bbox"][0] for component in components),
254 min(component["bbox"][1] for component in components),
255 max(component["bbox"][2] for component in components),
256 max(component["bbox"][3] for component in components),
257 )
258
259
260def extract_slot_frames(strip: Image.Image, frame_count: int) -> list[Image.Image]:
261 slot_width = strip.width / frame_count
262 frames = []
263 for index in range(frame_count):
264 left = round(index * slot_width)
265 right = round((index + 1) * slot_width)
266 crop = strip.crop((left, 0, right, strip.height))
267 frames.append(fit_to_cell(crop))
268 return frames
269
270
271def extract_stable_slot_frames(strip: Image.Image, frame_count: int) -> list[Image.Image]:
272 groups = component_frame_groups(strip, frame_count)
273 padding = 4
274 if groups is not None:
275 bboxes = [component_bounds(group) for group in groups]
276 shared_top = max(0, min(bbox[1] for bbox in bboxes) - padding)
277 shared_bottom = min(strip.height, max(bbox[3] for bbox in bboxes) + padding)
278 viewport_width = max(bbox[2] - bbox[0] for bbox in bboxes) + padding * 2
279 viewport_height = max(1, shared_bottom - shared_top)
280 frames = []
281 for group, bbox in zip(groups, bboxes):
282 grouped = component_group_image(strip, group, padding=padding)
283 grouped_top = max(0, bbox[1] - padding)
284 viewport = Image.new(
285 "RGBA",
286 (viewport_width, viewport_height),
287 (0, 0, 0, 0),
288 )
289 left = (viewport_width - grouped.width) // 2
290 viewport.alpha_composite(grouped, (left, grouped_top - shared_top))
291 frames.append(fit_viewport_to_cell(viewport))
292 return frames
293
294 bbox = strip.getbbox()
295 if bbox is None:
296 return [
297 Image.new("RGBA", (CELL_WIDTH, CELL_HEIGHT), (0, 0, 0, 0))
298 for _ in range(frame_count)
299 ]
300
301 shared_top = max(0, bbox[1] - padding)
302 shared_bottom = min(strip.height, bbox[3] + padding)
303 slot_width = strip.width / frame_count
304 frames = []
305 for index in range(frame_count):
306 left = round(index * slot_width)
307 right = round((index + 1) * slot_width)
308 crop = strip.crop((left, shared_top, right, shared_bottom))
309 frames.append(fit_viewport_to_cell(crop))
310 return frames
311
312
313def extract_state(
314 strip_path: Path,
315 state: str,
316 output_root: Path,
317 chroma_key: tuple[int, int, int],
318 threshold: float,
319 method: str,
320) -> dict[str, object]:
321 frame_count = ROW_FRAME_COUNTS[state]
322 with Image.open(strip_path) as opened:
323 strip = remove_chroma_background(opened, chroma_key, threshold)
324
325 state_dir = output_root / state
326 state_dir.mkdir(parents=True, exist_ok=True)
327
328 frames = None
329 used_method = method
330 if method in {"auto", "components"}:
331 frames = extract_component_frames(strip, frame_count)
332 if frames is None and method == "components":
333 raise SystemExit(f"could not find {frame_count} sprite components in {strip_path}")
334 if frames is not None:
335 used_method = "components"
336
337 if frames is None:
338 if method == "stable-slots":
339 frames = extract_stable_slot_frames(strip, frame_count)
340 used_method = "stable-slots"
341 else:
342 frames = extract_slot_frames(strip, frame_count)
343 used_method = "slots"
344
345 outputs = []
346 for index, frame in enumerate(frames):
347 output = state_dir / f"{index:02d}.png"
348 frame.save(output)
349 outputs.append(str(output))
350 return {"state": state, "frames": outputs, "method": used_method}
351
352
353def main() -> None:
354 parser = argparse.ArgumentParser(description=__doc__)
355 parser.add_argument("--decoded-dir", required=True)
356 parser.add_argument("--output-dir", required=True)
357 parser.add_argument("--states", default="all")
358 parser.add_argument("--chroma-key", help="Override chroma key as #RRGGBB.")
359 parser.add_argument("--key-threshold", type=float, default=96.0)
360 parser.add_argument(
361 "--method",
362 choices=("auto", "components", "slots", "stable-slots"),
363 default="auto",
364 help="Use connected sprite components when possible, raw equal slots, or row-stable slot viewports.",
365 )
366 args = parser.parse_args()
367
368 decoded_dir = Path(args.decoded_dir).expanduser().resolve()
369 output_dir = Path(args.output_dir).expanduser().resolve()
370 chroma_key = load_chroma_key(decoded_dir, args.chroma_key)
371 states = parse_states(args.states)
372 manifest = []
373 for state in states:
374 strip_path = decoded_dir / f"{state}.png"
375 if not strip_path.is_file():
376 raise SystemExit(f"missing generated strip for {state}: {strip_path}")
377 manifest.append(
378 extract_state(
379 strip_path,
380 state,
381 output_dir,
382 chroma_key,
383 args.key_threshold,
384 args.method,
385 )
386 )
387
388 (output_dir / "frames-manifest.json").write_text(
389 json.dumps(
390 {
391 "ok": True,
392 "chroma_key": {
393 "hex": f"#{chroma_key[0]:02X}{chroma_key[1]:02X}{chroma_key[2]:02X}",
394 "rgb": list(chroma_key),
395 "threshold": args.key_threshold,
396 },
397 "rows": manifest,
398 },
399 indent=2,
400 )
401 + "\n",
402 encoding="utf-8",
403 )
404 print(json.dumps({"ok": True, "frames_root": str(output_dir), "states": states}, indent=2))
405
406
407if __name__ == "__main__":
408 main()