from __future__ import annotations import sys from dataclasses import dataclass from pathlib import Path from PIL import Image, ImageChops, ImageEnhance, ImageFilter, ImageOps FRAME_SIZE = 313 GRID_SIZE = 4 PALETTE_COLORS = 192 ROOT = Path(__file__).resolve().parents[1] UNIT_DIR = ROOT / "src" / "assets" / "images" / "units" REFERENCE_SHEET = UNIT_DIR / "unit-cao-cao.png" @dataclass(frozen=True) class FrameTarget: center_x: float bottom_y: int width: int height: int @dataclass(frozen=True) class UnitStyle: scale: float saturation: float contrast: float brightness: float edge_warmth: float rim_color: tuple[int, int, int] @dataclass(frozen=True) class TemplateVariant: prefix: str template: str tint: tuple[int, int, int] strength: float saturation: float = 1.0 brightness: float = 1.0 TEMPLATE_VARIANTS = ( TemplateVariant("unit-shu-archer", "unit-rebel-archer", (52, 108, 74), 0.18, 0.88, 0.98), TemplateVariant("unit-shu-cavalry", "unit-ma-chao", (54, 104, 72), 0.16, 0.86, 0.98), TemplateVariant("unit-shu-infantry", "unit-rebel", (54, 108, 72), 0.2, 0.88, 0.98), TemplateVariant("unit-shu-officer", "unit-guan-yu", (56, 108, 76), 0.14, 0.86, 0.98), TemplateVariant("unit-shu-spearman", "unit-zhang-fei", (68, 104, 82), 0.13, 0.86, 0.98), TemplateVariant("unit-shu-strategist", "unit-zhuge-liang", (70, 112, 92), 0.14, 0.86, 0.99), TemplateVariant("unit-wei-archer", "unit-rebel-archer", (68, 76, 104), 0.2, 0.86, 0.96), TemplateVariant("unit-wei-cavalry", "unit-ma-chao", (70, 72, 102), 0.2, 0.84, 0.96), TemplateVariant("unit-wei-infantry", "unit-rebel", (66, 72, 100), 0.22, 0.84, 0.96), TemplateVariant("unit-wei-officer", "unit-cao-cao", (74, 70, 104), 0.16, 0.84, 0.96), TemplateVariant("unit-wei-strategist", "unit-zhuge-liang", (74, 84, 116), 0.18, 0.84, 0.96), TemplateVariant("unit-wu-archer", "unit-rebel-archer", (42, 98, 104), 0.2, 0.88, 0.98), TemplateVariant("unit-wu-cavalry", "unit-ma-chao", (42, 96, 104), 0.18, 0.86, 0.98), TemplateVariant("unit-wu-infantry", "unit-rebel", (42, 100, 104), 0.22, 0.88, 0.98), TemplateVariant("unit-wu-officer", "unit-lu-meng", (46, 100, 110), 0.16, 0.86, 0.98), TemplateVariant("unit-wu-strategist", "unit-zhuge-liang", (44, 102, 112), 0.18, 0.86, 0.98), TemplateVariant("unit-nanman-infantry", "unit-rebel", (104, 96, 48), 0.22, 0.9, 0.98), TemplateVariant("unit-nanman-officer", "unit-meng-huo", (108, 88, 50), 0.17, 0.88, 0.98), TemplateVariant("unit-nanman-shaman", "unit-zhuge-liang", (104, 98, 52), 0.2, 0.88, 0.98), ) def main() -> None: generic_only = "--generic-only" in sys.argv cleanup_only = "--cleanup-only" in sys.argv optimize_only = "--optimize-only" in sys.argv stabilize_motion = "--stabilize-motion" in sys.argv tone_only = "--tone-only" in sys.argv selected_modes = sum(1 for enabled in (generic_only, cleanup_only, optimize_only, stabilize_motion, tone_only) if enabled) if selected_modes > 1: raise ValueError("Use only one of --generic-only, --cleanup-only, --optimize-only, --stabilize-motion, or --tone-only") if generic_only: sheets = load_current_base_sheets() write_template_variants(sheets) return if cleanup_only: cleanup_current_base_sheets() return if optimize_only: optimize_current_unit_pngs() return if stabilize_motion: stabilize_current_motion_sheets() return if tone_only: tone_current_unit_pngs() return reference = Image.open(REFERENCE_SHEET).convert("RGBA") targets = reference_targets(reference) generated = 0 generated_sheets: dict[str, Image.Image] = {} for source in sorted(UNIT_DIR.glob("unit-*.png")): if source.name.endswith("-actions.png"): continue sheet = Image.open(source).convert("RGBA") if sheet.size != (FRAME_SIZE * GRID_SIZE, FRAME_SIZE * GRID_SIZE): raise ValueError(f"{source.name} has unexpected size {sheet.size}") if source.name == REFERENCE_SHEET.name: output = ensure_reference_finish(sheet) else: output = harmonize_sheet(sheet, source.stem, targets) save_unit_png(output, source) generated_sheets[source.stem] = output generated += 1 write_template_variants(generated_sheets) print(f"Harmonized {generated} unit sheets against {REFERENCE_SHEET.name}") def load_current_base_sheets() -> dict[str, Image.Image]: sheets: dict[str, Image.Image] = {} for source in sorted(UNIT_DIR.glob("unit-*.png")): if source.name.endswith("-actions.png"): continue sheets[source.stem] = Image.open(source).convert("RGBA") return sheets def cleanup_current_base_sheets() -> None: cleaned = 0 for source in sorted(UNIT_DIR.glob("unit-*.png")): if source.name.endswith("-actions.png"): continue sheet = Image.open(source).convert("RGBA") output = Image.new("RGBA", sheet.size, (0, 0, 0, 0)) for row in range(GRID_SIZE): for col in range(GRID_SIZE): frame = crop_frame(sheet, row, col) output.alpha_composite(remove_detached_artifacts(clean_alpha(frame)), (col * FRAME_SIZE, row * FRAME_SIZE)) save_unit_png(output, source) cleaned += 1 print(f"Cleaned detached artifacts in {cleaned} unit sheets") def optimize_current_unit_pngs() -> None: optimized = 0 for source in sorted(UNIT_DIR.glob("unit-*.png")): image = Image.open(source).convert("RGBA") save_unit_png(image, source) optimized += 1 print(f"Optimized {optimized} unit PNGs with {PALETTE_COLORS}-color palettes") def tone_current_unit_pngs() -> None: toned = 0 for source in sorted(UNIT_DIR.glob("unit-*.png")): image = Image.open(source).convert("RGBA") save_unit_png(stabilize_palette_image(image), source) toned += 1 print(f"Applied muted battlefield tone to {toned} unit PNGs") def stabilize_current_motion_sheets() -> None: reference = Image.open(REFERENCE_SHEET).convert("RGBA") targets = stable_reference_targets(reference_targets(reference)) reference_east_score = side_orientation_score(crop_frame(reference, 1, 0)) generated_sheets: dict[str, Image.Image] = {} stabilized = 0 for source in sorted(UNIT_DIR.glob("unit-*.png")): if source.name.endswith("-actions.png"): continue sheet = Image.open(source).convert("RGBA") if sheet.size != (FRAME_SIZE * GRID_SIZE, FRAME_SIZE * GRID_SIZE): raise ValueError(f"{source.name} has unexpected size {sheet.size}") output = stabilize_motion_sheet(sheet, source.stem, targets, reference_east_score) save_unit_png(output, source) generated_sheets[source.stem] = output stabilized += 1 write_template_variants(generated_sheets) print(f"Stabilized directional motion in {stabilized} unit sheets") def stable_reference_targets(targets: list[list[FrameTarget]]) -> list[list[FrameTarget]]: stable_rows: list[list[FrameTarget]] = [] for row_targets in targets: width = round(sum(target.width for target in row_targets) / len(row_targets)) height = round(sum(target.height for target in row_targets) / len(row_targets)) stable_rows.append( [ FrameTarget( center_x=target.center_x, bottom_y=target.bottom_y, width=width, height=height, ) for target in row_targets ] ) return stable_rows def stabilize_motion_sheet( sheet: Image.Image, stem: str, targets: list[list[FrameTarget]], reference_east_score: float, ) -> Image.Image: style = style_for_unit(stem) sources = { 0: canonical_row_frame(sheet, 0), 1: oriented_east_frame(canonical_row_frame(sheet, 1), reference_east_score), 2: canonical_row_frame(sheet, 2), } sources[3] = ImageOps.mirror(sources[1]) output = Image.new("RGBA", sheet.size, (0, 0, 0, 0)) for row in range(GRID_SIZE): source = sources[row] for col in range(GRID_SIZE): frame = normalize_frame(source, targets[row][col], style) output.alpha_composite(frame, (col * FRAME_SIZE, row * FRAME_SIZE)) return output def canonical_row_frame(sheet: Image.Image, row: int) -> Image.Image: preferred = remove_detached_artifacts(clean_alpha(crop_frame(sheet, row, 0))) preferred_area = alpha_area(preferred) candidates = [remove_detached_artifacts(clean_alpha(crop_frame(sheet, row, col))) for col in range(GRID_SIZE)] best = max(candidates, key=alpha_area) best_area = alpha_area(best) if preferred_area >= max(1, best_area) * 0.55: return preferred return best def oriented_east_frame(frame: Image.Image, reference_east_score: float) -> Image.Image: score = side_orientation_score(frame) if abs(score) < 1.25 or abs(reference_east_score) < 1.25: return frame return ImageOps.mirror(frame) if score * reference_east_score < 0 else frame def alpha_area(image: Image.Image) -> int: alpha = image.getchannel("A") histogram = alpha.histogram() return sum(count for value, count in enumerate(histogram) if value > 12) def side_orientation_score(image: Image.Image) -> float: image = remove_detached_artifacts(clean_alpha(image)) bbox = alpha_bbox(image) if not bbox: return 0.0 pixels = image.load() left, top, right, bottom = bbox center_x = (left + right) / 2 upper_bottom = top + round((bottom - top) * 0.68) weighted_offset = 0.0 weight_total = 0.0 for y in range(top, upper_bottom): for x in range(left, right): red, green, blue, alpha = pixels[x, y] if alpha <= 24: continue brightness = (red + green + blue) / 3 warmth = red * 1.1 + green * 0.45 - blue * 0.45 weight = max(0.0, brightness - 38) + max(0.0, warmth - 58) * 0.34 if weight <= 0: continue weight *= alpha / 255 weighted_offset += (x - center_x) * weight weight_total += weight return weighted_offset / weight_total if weight_total else 0.0 def write_template_variants(sheets: dict[str, Image.Image]) -> None: written = 0 for source in sorted(UNIT_DIR.glob("unit-*.png")): if source.name.endswith("-actions.png"): continue variant = template_variant_for(source.stem) if not variant: continue template = sheets.get(variant.template) if not template: raise ValueError(f"Missing template sheet {variant.template} for {source.stem}") output = tint_template(template, variant) save_unit_png(output, source) sheets[source.stem] = output written += 1 print(f"Wrote {written} high-detail generic unit variants") def save_unit_png(image: Image.Image, path: Path) -> None: image = stabilize_palette_image(image.convert("RGBA")) indexed = image.quantize(colors=PALETTE_COLORS, method=Image.Quantize.FASTOCTREE, dither=Image.Dither.NONE) indexed.save(path, optimize=True) def stabilize_palette_image(image: Image.Image) -> Image.Image: image = image.convert("RGBA") red, green, blue, alpha = image.split() rgb = Image.merge("RGB", (red, green, blue)) rgb = ImageEnhance.Contrast(rgb).enhance(1.01) rgb = ImageEnhance.Brightness(rgb).enhance(1.0) return Image.merge("RGBA", (*rgb.split(), alpha)) def template_variant_for(stem: str) -> TemplateVariant | None: for variant in TEMPLATE_VARIANTS: if stem.startswith(variant.prefix) and stem != variant.template: return variant return None def tint_template(sheet: Image.Image, variant: TemplateVariant) -> Image.Image: rgb = sheet.convert("RGB") alpha = sheet.getchannel("A") luma = ImageOps.grayscale(rgb) dark = tuple(max(0, round(channel * 0.2)) for channel in variant.tint) light = tuple(min(255, round(channel * 1.18 + 84)) for channel in variant.tint) colorized = ImageOps.colorize(luma, black=dark, white=light, mid=variant.tint) mask = alpha.point(lambda value: min(255, round(value * variant.strength))) mixed = Image.composite(colorized, rgb, mask) mixed = ImageEnhance.Color(mixed).enhance(variant.saturation) mixed = ImageEnhance.Brightness(mixed).enhance(variant.brightness) return Image.merge("RGBA", (*mixed.split(), alpha)) def reference_targets(reference: Image.Image) -> list[list[FrameTarget]]: rows: list[list[FrameTarget]] = [] for row in range(GRID_SIZE): target_row: list[FrameTarget] = [] for col in range(GRID_SIZE): frame = crop_frame(reference, row, col) bbox = alpha_bbox(frame) if not bbox: target_row.append(FrameTarget(FRAME_SIZE / 2, FRAME_SIZE - 8, 172, 282)) continue left, top, right, bottom = bbox target_row.append( FrameTarget( center_x=(left + right) / 2, bottom_y=min(FRAME_SIZE - 2, bottom), width=right - left, height=bottom - top, ) ) rows.append(target_row) return rows def harmonize_sheet(sheet: Image.Image, stem: str, targets: list[list[FrameTarget]]) -> Image.Image: output = Image.new("RGBA", sheet.size, (0, 0, 0, 0)) style = style_for_unit(stem) for row in range(GRID_SIZE): for col in range(GRID_SIZE): frame = crop_frame(sheet, row, col) target = targets[row][col] normalized = normalize_frame(frame, target, style) output.alpha_composite(normalized, (col * FRAME_SIZE, row * FRAME_SIZE)) return output def crop_frame(sheet: Image.Image, row: int, col: int) -> Image.Image: return sheet.crop((col * FRAME_SIZE, row * FRAME_SIZE, (col + 1) * FRAME_SIZE, (row + 1) * FRAME_SIZE)).copy() def normalize_frame(frame: Image.Image, target: FrameTarget, style: UnitStyle) -> Image.Image: frame = remove_detached_artifacts(clean_alpha(frame)) bbox = alpha_bbox(frame) if not bbox: return frame subject = frame.crop(bbox) width, height = subject.size fit_width = FRAME_SIZE - 16 fit_height = FRAME_SIZE - 8 target_height = min(fit_height, max(218, target.height * style.scale)) target_width = min(fit_width, max(132, target.width * style.scale)) scale = 1.0 if height > target_height * 1.06 or width > fit_width: scale = min(target_height / height, fit_width / width) elif height < target_height * 0.9: scale = min(target_height / height, 1.13) if width * scale > target_width * 1.22: scale = min(scale, (target_width * 1.22) / width) new_width = max(1, min(fit_width, round(width * scale))) new_height = max(1, min(fit_height, round(height * scale))) if (new_width, new_height) != subject.size: subject = subject.resize((new_width, new_height), Image.Resampling.LANCZOS) subject = grade_subject(subject, style) canvas = Image.new("RGBA", frame.size, (0, 0, 0, 0)) left = round(target.center_x - new_width / 2) top = round(target.bottom_y - new_height) left = max(4, min(FRAME_SIZE - new_width - 4, left)) top = max(0, min(FRAME_SIZE - new_height - 2, top)) canvas.alpha_composite(subject, (left, top)) return finish_frame(canvas, style) def ensure_reference_finish(sheet: Image.Image) -> Image.Image: style = style_for_unit("unit-cao-cao") output = Image.new("RGBA", sheet.size, (0, 0, 0, 0)) for row in range(GRID_SIZE): for col in range(GRID_SIZE): frame = crop_frame(sheet, row, col) output.alpha_composite(finish_frame(remove_detached_artifacts(clean_alpha(frame)), style), (col * FRAME_SIZE, row * FRAME_SIZE)) return output def clean_alpha(image: Image.Image) -> Image.Image: red, green, blue, alpha = image.split() alpha = alpha.point(lambda value: 0 if value < 7 else (255 if value > 248 else value)) return Image.merge("RGBA", (red, green, blue, alpha)) def remove_detached_artifacts(image: Image.Image) -> Image.Image: alpha = image.getchannel("A") pixels = alpha.load() width, height = alpha.size visited = bytearray(width * height) components: list[tuple[int, tuple[int, int, int, int], list[tuple[int, int]]]] = [] for y in range(height): for x in range(width): index = y * width + x if visited[index] or pixels[x, y] <= 12: continue stack = [(x, y)] visited[index] = 1 points: list[tuple[int, int]] = [] left = right = x top = bottom = y while stack: px, py = stack.pop() points.append((px, py)) left = min(left, px) right = max(right, px) top = min(top, py) bottom = max(bottom, py) for nx, ny in ((px + 1, py), (px - 1, py), (px, py + 1), (px, py - 1)): if nx < 0 or nx >= width or ny < 0 or ny >= height: continue neighbor_index = ny * width + nx if visited[neighbor_index] or pixels[nx, ny] <= 12: continue visited[neighbor_index] = 1 stack.append((nx, ny)) components.append((len(points), (left, top, right + 1, bottom + 1), points)) if len(components) <= 1: return image components.sort(key=lambda component: component[0], reverse=True) largest_area, largest_bbox, _ = components[0] keep_points: set[tuple[int, int]] = set(components[0][2]) for area, bbox, points in components[1:]: distance = bbox_distance(largest_bbox, bbox) if distance <= 3 or area >= largest_area * 0.18: keep_points.update(points) cleaned_alpha = Image.new("L", image.size, 0) cleaned_pixels = cleaned_alpha.load() source_pixels = alpha.load() for x, y in keep_points: cleaned_pixels[x, y] = source_pixels[x, y] red, green, blue, _ = image.split() return Image.merge("RGBA", (red, green, blue, cleaned_alpha)) def bbox_distance(left_box: tuple[int, int, int, int], right_box: tuple[int, int, int, int]) -> int: left_a, top_a, right_a, bottom_a = left_box left_b, top_b, right_b, bottom_b = right_box dx = max(left_a - right_b, left_b - right_a, 0) dy = max(top_a - bottom_b, top_b - bottom_a, 0) return max(dx, dy) def alpha_bbox(image: Image.Image) -> tuple[int, int, int, int] | None: alpha = image.getchannel("A") return alpha.getbbox() def grade_subject(subject: Image.Image, style: UnitStyle) -> Image.Image: red, green, blue, alpha = subject.split() rgb = Image.merge("RGB", (red, green, blue)) rgb = ImageEnhance.Color(rgb).enhance(style.saturation) rgb = ImageEnhance.Contrast(rgb).enhance(style.contrast) rgb = ImageEnhance.Brightness(rgb).enhance(style.brightness) poster = ImageOps.posterize(rgb, 6) rgb = Image.blend(rgb, poster, 0.18) rgb = rgb.filter(ImageFilter.UnsharpMask(radius=1.0, percent=120, threshold=3)) return Image.merge("RGBA", (*rgb.split(), alpha)) def finish_frame(frame: Image.Image, style: UnitStyle) -> Image.Image: alpha = frame.getchannel("A") outer_alpha = ImageChops.subtract(alpha.filter(ImageFilter.MaxFilter(7)), alpha) ink_alpha = ImageChops.subtract(alpha.filter(ImageFilter.MaxFilter(3)), alpha) glow_alpha = alpha.filter(ImageFilter.GaussianBlur(2.2)) inner_alpha = ImageChops.subtract(alpha, alpha.filter(ImageFilter.MinFilter(3))) result = Image.new("RGBA", frame.size, (0, 0, 0, 0)) result.alpha_composite(colored_mask(glow_alpha, (*style.rim_color, 34))) result.alpha_composite(colored_mask(outer_alpha, (*style.rim_color, 160))) result.alpha_composite(colored_mask(ink_alpha, (10, 8, 13, 185))) warm = colored_mask(inner_alpha, (214, 156, 58, round(34 * style.edge_warmth))) result.alpha_composite(frame) result.alpha_composite(warm) return clean_alpha(result) def colored_mask(alpha: Image.Image, color: tuple[int, int, int, int]) -> Image.Image: red, green, blue, opacity = color mask = alpha.point(lambda value: min(255, round(value * opacity / 255))) layer = Image.new("RGBA", alpha.size, (red, green, blue, 0)) layer.putalpha(mask) return layer def style_for_unit(stem: str) -> UnitStyle: lower = stem.lower() scale = 0.98 saturation = 1.0 contrast = 1.04 brightness = 1.0 edge_warmth = 0.9 rim_color = (31, 25, 35) if any(token in lower for token in ("archer", "crossbow", "longbow")): scale = 0.94 saturation = 1.0 elif any(token in lower for token in ("cavalry", "ma-chao", "ma-dai", "lu-bu")): scale = 1.03 contrast = 1.05 elif any(token in lower for token in ("strategist", "shaman", "zhuge", "sima", "fa-zheng", "ma-liang", "yi-ji")): scale = 0.93 brightness = 1.0 rim_color = (29, 28, 41) elif any(token in lower for token in ("leader", "officer", "meng-huo", "guan-yu", "zhang-fei", "jiang-wei")): scale = 1.01 contrast = 1.05 edge_warmth = 0.98 if "nanman" in lower or "meng-huo" in lower: saturation = 1.0 rim_color = (38, 28, 34) elif "wu-" in lower or "lu-meng" in lower: rim_color = (25, 31, 37) elif "wei-" in lower or "cao-cao" in lower or "sima" in lower: rim_color = (30, 28, 40) elif "shu-" in lower or any(token in lower for token in ("liu-bei", "guan-yu", "zhang-fei", "zhao-yun")): rim_color = (28, 30, 38) return UnitStyle(scale, saturation, contrast, brightness, edge_warmth, rim_color) if __name__ == "__main__": main()