feat: ui rework and gear generator

chainring_generator/generate_chainrings.py

@@ -0,0 +1,492 @@
+from __future__ import annotations
+
+import argparse
+import csv
+import math
+import shutil
+import subprocess
+from dataclasses import dataclass
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+import numpy as np
+from PIL import Image, ImageOps
+
+
+@dataclass(frozen=True)
+class Canvas:
+    width: int = 588
+    height: int = 495
+    cx: float = 304.0
+    cy: float = 256.0
+
+
+@dataclass(frozen=True)
+class Style:
+    background: str = "#3d3d3d"
+    material: str = "#eeeeee"
+    outline: str = "#f4f4f4"
+    outline_width: float = 3.0
+
+
+@dataclass(frozen=True)
+class Geometry:
+    root_radius: float = 213.0
+    tooth_depth: float = 17.5
+    reference_teeth: int = 60
+    pitch_radius_offset: float = 8.0
+    tooth_sharpness: float = 1.32
+    low_tooth_threshold: int = 25
+    low_tooth_half_width: float = 0.30
+    low_tooth_cusp: float = 0.55
+    central_hole_radius: float = 46.0
+    bolt_hole_radius: float = 16.5
+    material_margin: float = 8.0
+    bolt_hole_min_teeth: int = 25
+    decorative_cutout_min_scale: float = 0.62
+    samples_per_tooth: int = 28
+
+
+@dataclass(frozen=True)
+class SprocketSizing:
+    pitch_radius: float
+    calculated_root_radius: float
+    root_radius: float
+    tooth_depth: float
+    outer_radius: float
+    decorative_scale: float
+    bolt_scale: float
+    include_bolt_holes: bool
+    include_decorative_cutouts: bool
+
+
+@dataclass(frozen=True)
+class Cutout:
+    kind: str
+    cx: float
+    cy: float
+    rx: float
+    ry: float | None = None
+    rotation: float = 0.0
+
+
+CANVAS = Canvas()
+STYLE = Style()
+GEOMETRY = Geometry()
+
+
+def fmt(value: float) -> str:
+    return f"{value:.3f}".rstrip("0").rstrip(".")
+
+
+def polygon_path(points: list[tuple[float, float]]) -> str:
+    start = points[0]
+    rest = points[1:]
+    commands = [f"M {fmt(start[0])} {fmt(start[1])}"]
+    commands.extend(f"L {fmt(x)} {fmt(y)}" for x, y in rest)
+    commands.append("Z")
+    return " ".join(commands)
+
+
+def circle_points(cx: float, cy: float, radius: float, samples: int = 144) -> list[tuple[float, float]]:
+    return [
+        (cx + math.cos(t) * radius, cy + math.sin(t) * radius)
+        for t in np.linspace(0.0, 2.0 * math.pi, samples, endpoint=False)
+    ]
+
+
+def ellipse_points(
+    cx: float,
+    cy: float,
+    rx: float,
+    ry: float,
+    rotation_deg: float = 0.0,
+    samples: int = 192,
+) -> list[tuple[float, float]]:
+    rotation = math.radians(rotation_deg)
+    cos_r = math.cos(rotation)
+    sin_r = math.sin(rotation)
+    points = []
+    for t in np.linspace(0.0, 2.0 * math.pi, samples, endpoint=False):
+        x = math.cos(t) * rx
+        y = math.sin(t) * ry
+        points.append((cx + x * cos_r - y * sin_r, cy + x * sin_r + y * cos_r))
+    return points
+
+
+DECORATIVE_CUTOUT_SPECS = (
+    ("ellipse", 0.0, -141.0, 121.0, 41.0, 0.0),
+    ("ellipse", -109.0, 58.0, 106.0, 41.0, 63.0),
+    ("ellipse", 109.0, 58.0, 106.0, 41.0, -63.0),
+)
+
+BOLT_HOLE_OFFSETS = (
+    (-134.0, -76.0),
+    (134.0, -76.0),
+    (0.0, 157.0),
+)
+
+REFERENCE_BOLT_OFFSET_RADIUS = max(math.hypot(dx, dy) for dx, dy in BOLT_HOLE_OFFSETS)
+
+
+def decorative_reference_extent() -> float:
+    extent = 0.0
+    for _, dx, dy, rx, ry, rotation in DECORATIVE_CUTOUT_SPECS:
+        points = ellipse_points(dx, dy, rx, ry, rotation, 240)
+        extent = max(extent, *(math.hypot(x, y) for x, y in points))
+    return extent
+
+
+DECORATIVE_REFERENCE_EXTENT = decorative_reference_extent()
+
+
+def clamp(value: float, minimum: float, maximum: float) -> float:
+    return max(minimum, min(maximum, value))
+
+
+def calculated_pitch_radius(teeth: int, geometry: Geometry = GEOMETRY) -> float:
+    reference_pitch_radius = geometry.root_radius + geometry.pitch_radius_offset
+    return reference_pitch_radius * math.sin(math.pi / geometry.reference_teeth) / math.sin(math.pi / teeth)
+
+
+def sprocket_sizing(teeth: int, geometry: Geometry = GEOMETRY) -> SprocketSizing:
+    pitch_radius = calculated_pitch_radius(teeth, geometry)
+    calculated_root_radius = pitch_radius - geometry.pitch_radius_offset
+    central_min_root_radius = geometry.central_hole_radius + geometry.material_margin
+
+    include_bolt_holes = teeth >= geometry.bolt_hole_min_teeth
+    if include_bolt_holes:
+        bolt_min_radius = geometry.central_hole_radius + geometry.bolt_hole_radius + geometry.material_margin
+        bolt_min_root_radius = bolt_min_radius + geometry.bolt_hole_radius + geometry.material_margin
+        min_root_radius = max(central_min_root_radius, bolt_min_root_radius)
+    else:
+        min_root_radius = central_min_root_radius
+
+    root_radius = max(calculated_root_radius, min_root_radius)
+    outer_radius = root_radius + geometry.tooth_depth
+
+    max_bolt_offset = max(0.0, root_radius - geometry.bolt_hole_radius - geometry.material_margin)
+    if include_bolt_holes:
+        min_bolt_offset = geometry.central_hole_radius + geometry.bolt_hole_radius + geometry.material_margin
+        bolt_offset = clamp(max_bolt_offset, min_bolt_offset, REFERENCE_BOLT_OFFSET_RADIUS)
+        bolt_scale = bolt_offset / REFERENCE_BOLT_OFFSET_RADIUS
+    else:
+        bolt_scale = 0.0
+
+    decorative_scale = clamp((root_radius - geometry.material_margin) / DECORATIVE_REFERENCE_EXTENT, 0.0, 1.0)
+    include_decorative_cutouts = decorative_scale >= geometry.decorative_cutout_min_scale
+
+    return SprocketSizing(
+        pitch_radius=pitch_radius,
+        calculated_root_radius=calculated_root_radius,
+        root_radius=root_radius,
+        tooth_depth=geometry.tooth_depth,
+        outer_radius=outer_radius,
+        decorative_scale=decorative_scale,
+        bolt_scale=bolt_scale,
+        include_bolt_holes=include_bolt_holes,
+        include_decorative_cutouts=include_decorative_cutouts,
+    )
+
+
+def outer_points(
+    teeth: int,
+    geometry: Geometry = GEOMETRY,
+    canvas: Canvas = CANVAS,
+    radius_delta: float = 0.0,
+) -> list[tuple[float, float]]:
+    samples = teeth * geometry.samples_per_tooth
+    sizing = sprocket_sizing(teeth, geometry)
+    points = []
+    phase_offset = -math.pi / 2.0
+
+    for index in range(samples):
+        theta = phase_offset + 2.0 * math.pi * index / samples
+        tooth_phase = (index % geometry.samples_per_tooth) / geometry.samples_per_tooth
+        radius = sizing.root_radius + sizing.tooth_depth * tooth_profile(teeth, tooth_phase, geometry) + radius_delta
+        points.append((canvas.cx + math.cos(theta) * radius, canvas.cy + math.sin(theta) * radius))
+
+    return points
+
+
+def tooth_profile(teeth: int, phase: float, geometry: Geometry = GEOMETRY) -> float:
+    if teeth < geometry.low_tooth_threshold:
+        distance_from_peak = abs(phase - 0.5)
+        if distance_from_peak >= geometry.low_tooth_half_width:
+            return 0.0
+        return 1.0 - (distance_from_peak / geometry.low_tooth_half_width) ** geometry.low_tooth_cusp
+
+    smooth_profile = 0.5 - 0.5 * math.cos(2.0 * math.pi * phase)
+    return smooth_profile**geometry.tooth_sharpness
+
+
+def offset_cutout(
+    kind: str,
+    dx: float,
+    dy: float,
+    rx: float,
+    ry: float | None,
+    rotation: float,
+    center_scale: float,
+    radius_scale: float,
+    canvas: Canvas = CANVAS,
+) -> Cutout:
+    scaled_ry = None if ry is None else ry * radius_scale
+    return Cutout(
+        kind,
+        canvas.cx + dx * center_scale,
+        canvas.cy + dy * center_scale,
+        rx * radius_scale,
+        scaled_ry,
+        rotation,
+    )
+
+
+def cutouts_for_teeth(teeth: int, canvas: Canvas = CANVAS, geometry: Geometry = GEOMETRY) -> list[Cutout]:
+    sizing = sprocket_sizing(teeth, geometry)
+    return [
+        *(
+            offset_cutout(kind, dx, dy, rx, ry, rotation, sizing.decorative_scale, sizing.decorative_scale, canvas)
+            for kind, dx, dy, rx, ry, rotation in DECORATIVE_CUTOUT_SPECS
+            if sizing.include_decorative_cutouts
+        ),
+        offset_cutout("circle", 0.0, 0.0, geometry.central_hole_radius, None, 0.0, 1.0, 1.0, canvas),
+        *(
+            offset_cutout("circle", dx, dy, geometry.bolt_hole_radius, None, 0.0, sizing.bolt_scale, 1.0, canvas)
+            for dx, dy in BOLT_HOLE_OFFSETS
+            if sizing.include_bolt_holes
+        ),
+    ]
+
+
+def cutout_path(cutout: Cutout, delta: float = 0.0, reverse: bool = False) -> str:
+    if cutout.kind == "circle":
+        points = circle_points(cutout.cx, cutout.cy, cutout.rx + delta, 160)
+    elif cutout.kind == "ellipse":
+        assert cutout.ry is not None
+        points = ellipse_points(
+            cutout.cx,
+            cutout.cy,
+            cutout.rx + delta,
+            cutout.ry + delta,
+            cutout.rotation,
+            240,
+        )
+    else:
+        raise ValueError(f"Unsupported cutout kind: {cutout.kind}")
+
+    if reverse:
+        points = list(reversed(points))
+    return polygon_path(points)
+
+
+def cutout_paths_for_teeth(teeth: int, canvas: Canvas = CANVAS) -> list[str]:
+    return [cutout_path(cutout) for cutout in cutouts_for_teeth(teeth, canvas)]
+
+
+def filled_svg(teeth: int, canvas: Canvas = CANVAS, style: Style = STYLE) -> str:
+    outer = polygon_path(outer_points(teeth))
+    cutouts = " ".join(cutout_paths_for_teeth(teeth, canvas))
+    return f"""<svg xmlns="http://www.w3.org/2000/svg" width="{canvas.width}" height="{canvas.height}" viewBox="0 0 {canvas.width} {canvas.height}">
+  <path d="{outer} {cutouts}" fill="{style.material}" fill-rule="evenodd"/>
+</svg>
+"""
+
+
+def outline_svg(teeth: int, canvas: Canvas = CANVAS, style: Style = STYLE) -> str:
+    half_width = style.outline_width / 2.0
+    outer = polygon_path(outer_points(teeth, radius_delta=half_width))
+    inner = polygon_path(list(reversed(outer_points(teeth, radius_delta=-half_width))))
+    paths = [f"{outer} {inner}"]
+    for cutout in cutouts_for_teeth(teeth, canvas):
+        paths.append(f"{cutout_path(cutout, half_width)} {cutout_path(cutout, -half_width, reverse=True)}")
+
+    path_elements = "\n  ".join(
+        f'<path d="{path}" fill="{style.outline}" fill-rule="evenodd"/>'
+        for path in paths
+    )
+    return f"""<svg xmlns="http://www.w3.org/2000/svg" width="{canvas.width}" height="{canvas.height}" viewBox="0 0 {canvas.width} {canvas.height}">
+  {path_elements}
+</svg>
+"""
+
+
+def render_png(svg_path: Path, png_path: Path) -> None:
+    magick = shutil.which("magick") or shutil.which("convert")
+    if magick is None:
+        raise RuntimeError("ImageMagick is required for SVG-to-PNG conversion, but neither 'magick' nor 'convert' is on PATH.")
+
+    png_path.parent.mkdir(parents=True, exist_ok=True)
+    command = [
+        magick,
+        "-background",
+        "none",
+        "-density",
+        "192",
+        str(svg_path),
+        "-resize",
+        f"{CANVAS.width}x{CANVAS.height}!",
+        str(png_path),
+    ]
+    subprocess.run(command, check=True)
+
+
+def write_asset(teeth: int, variant: str, svg_root: Path, png_root: Path, png: bool = True) -> tuple[Path, Path | None]:
+    if variant == "filled":
+        svg_text = filled_svg(teeth)
+    elif variant == "outline":
+        svg_text = outline_svg(teeth)
+    else:
+        raise ValueError(f"Unsupported variant: {variant}")
+
+    stem = f"chainring_{teeth:02d}_{variant}"
+    svg_path = svg_root / variant / f"{stem}.svg"
+    png_path = png_root / variant / f"{stem}.png"
+    svg_path.parent.mkdir(parents=True, exist_ok=True)
+    svg_path.write_text(svg_text, encoding="utf-8")
+
+    if png:
+        render_png(svg_path, png_path)
+        return svg_path, png_path
+
+    return svg_path, None
+
+
+def generate_all(min_teeth: int, max_teeth: int, output_dir: Path, png: bool = True) -> list[Path]:
+    svg_root = output_dir / "svg"
+    png_root = output_dir / "png"
+    png_paths: list[Path] = []
+
+    for teeth in range(min_teeth, max_teeth + 1):
+        for variant in ("filled", "outline"):
+            _, png_path = write_asset(teeth, variant, svg_root, png_root, png=png)
+            if png_path is not None:
+                png_paths.append(png_path)
+
+    return png_paths
+
+
+def write_radii_manifest(min_teeth: int, max_teeth: int, output_path: Path) -> None:
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    fields = [
+        "teeth",
+        "pitch_radius_px",
+        "calculated_root_radius_px",
+        "root_radius_px",
+        "tooth_depth_px",
+        "outer_radius_px",
+        "bolt_scale",
+        "decorative_scale",
+        "bolt_holes",
+        "decorative_cutouts",
+    ]
+    with output_path.open("w", newline="", encoding="utf-8") as handle:
+        writer = csv.DictWriter(handle, fieldnames=fields)
+        writer.writeheader()
+        for teeth in range(min_teeth, max_teeth + 1):
+            sizing = sprocket_sizing(teeth)
+            writer.writerow(
+                {
+                    "teeth": teeth,
+                    "pitch_radius_px": f"{sizing.pitch_radius:.3f}",
+                    "calculated_root_radius_px": f"{sizing.calculated_root_radius:.3f}",
+                    "root_radius_px": f"{sizing.root_radius:.3f}",
+                    "tooth_depth_px": f"{sizing.tooth_depth:.3f}",
+                    "outer_radius_px": f"{sizing.outer_radius:.3f}",
+                    "bolt_scale": f"{sizing.bolt_scale:.3f}",
+                    "decorative_scale": f"{sizing.decorative_scale:.3f}",
+                    "bolt_holes": int(sizing.include_bolt_holes),
+                    "decorative_cutouts": int(sizing.include_decorative_cutouts),
+                }
+            )
+
+
+def make_contact_sheet(png_paths: list[Path], output_path: Path, columns: int = 6) -> None:
+    if not png_paths:
+        return
+
+    thumbs = []
+    thumb_size = (196, 165)
+    for path in png_paths:
+        image = Image.open(path).convert("RGBA")
+        image.thumbnail(thumb_size, Image.Resampling.LANCZOS)
+        tile = Image.new("RGBA", thumb_size, (61, 61, 61, 255))
+        x = (thumb_size[0] - image.width) // 2
+        y = (thumb_size[1] - image.height) // 2
+        tile.alpha_composite(image, (x, y))
+        thumbs.append(tile)
+
+    rows = math.ceil(len(thumbs) / columns)
+    sheet = Image.new("RGBA", (columns * thumb_size[0], rows * thumb_size[1]), (36, 36, 36, 255))
+    for index, thumb in enumerate(thumbs):
+        x = (index % columns) * thumb_size[0]
+        y = (index // columns) * thumb_size[1]
+        sheet.alpha_composite(thumb, (x, y))
+
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    sheet.save(output_path)
+
+
+def compare_to_reference(candidate_path: Path, reference_path: Path) -> dict[str, float]:
+    candidate = Image.open(candidate_path).convert("L")
+    reference = Image.open(reference_path).convert("L").resize(candidate.size, Image.Resampling.LANCZOS)
+
+    candidate_arr = np.asarray(ImageOps.autocontrast(candidate), dtype=np.float32) / 255.0
+    reference_arr = np.asarray(ImageOps.autocontrast(reference), dtype=np.float32) / 255.0
+    diff = candidate_arr - reference_arr
+
+    candidate_edges = np.abs(np.gradient(candidate_arr)[0]) + np.abs(np.gradient(candidate_arr)[1])
+    reference_edges = np.abs(np.gradient(reference_arr)[0]) + np.abs(np.gradient(reference_arr)[1])
+    edge_diff = candidate_edges - reference_edges
+
+    return {
+        "mse": float(np.mean(diff**2)),
+        "mae": float(np.mean(np.abs(diff))),
+        "edge_mae": float(np.mean(np.abs(edge_diff))),
+    }
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Generate chainring SVG and PNG assets with fixed reference-like inner cutouts.")
+    parser.add_argument("--min-teeth", type=int, default=5)
+    parser.add_argument("--max-teeth", type=int, default=60)
+    parser.add_argument("--output", type=Path, default=Path("out"))
+    parser.add_argument("--no-png", action="store_true")
+    parser.add_argument("--contact-sheet", action="store_true")
+    parser.add_argument("--reference", type=Path, help="Optional local reference image for reporting similarity metrics.")
+    parser.add_argument("--reference-teeth", type=int, default=44)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    if args.min_teeth < 3:
+        raise ValueError("--min-teeth must be at least 3.")
+    if args.max_teeth < args.min_teeth:
+        raise ValueError("--max-teeth must be greater than or equal to --min-teeth.")
+
+    png_paths = generate_all(args.min_teeth, args.max_teeth, args.output, png=not args.no_png)
+    write_radii_manifest(args.min_teeth, args.max_teeth, args.output / "radii.csv")
+
+    if args.contact_sheet and png_paths:
+        filled_samples = [path for path in png_paths if path.parent.name == "filled"]
+        outline_samples = [path for path in png_paths if path.parent.name == "outline"]
+        make_contact_sheet(filled_samples, args.output / "contact_sheet_filled.png")
+        make_contact_sheet(outline_samples, args.output / "contact_sheet_outline.png")
+
+    if args.reference and not args.no_png:
+        with TemporaryDirectory() as tmp:
+            tmp_root = Path(tmp)
+            _, candidate = write_asset(args.reference_teeth, "filled", tmp_root / "svg", tmp_root / "png", png=True)
+            assert candidate is not None
+            metrics = compare_to_reference(candidate, args.reference)
+        print(
+            "reference_metrics "
+            + " ".join(f"{name}={value:.6f}" for name, value in metrics.items())
+        )
+
+    print(f"generated teeth={args.min_teeth}..{args.max_teeth} output={args.output}")
+
+
+if __name__ == "__main__":
+    main()