iterative_refine.py

#!/usr/bin/env python3
"""
Iterative SVG Refinement - Progressive Vector Perfection

This script creates pixel-perfect SVG logos through multi-generation refinement:
1. Trace PNG → SVG (iteration 1)
2. Render SVG → PNG at high resolution
3. Enhance & smooth the PNG
4. Trace enhanced PNG → SVG (iteration 2)
5. Repeat for N iterations
6. Generate optimized versions at all required sizes

Each iteration progressively smooths and perfects the vector paths.
"""

import subprocess
import os
import sys
import tempfile
import shutil
from pathlib import Path
from PIL import Image, ImageFilter, ImageEnhance
from svg_optimizer import optimize_for_small_size
from rich.console import Console
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn

# Allow large images for iterative processing
Image.MAX_IMAGE_PIXELS = None

console = Console()


def ensure_inkscape():
    """Verify Inkscape is installed"""
    try:
        subprocess.run(["inkscape", "--version"], capture_output=True, check=True)
        return True
    except (subprocess.CalledProcessError, FileNotFoundError):
        console.print("[red]Error: Inkscape not found. Install with: brew install inkscape[/red]")
        return False


def enhance_image(input_path: str, output_path: str, iteration: int):
    """
    Enhance image quality for better tracing.
    Progressive enhancement - more aggressive in later iterations.
    """
    img = Image.open(input_path)

    # Convert to RGBA if needed
    if img.mode != 'RGBA':
        img = img.convert('RGBA')

    # Upscale for maximum detail capture, but cap at reasonable size
    # Only upscale on first iteration, then maintain size
    if iteration == 1:
        # First iteration: upscale to 2048px for maximum quality (monochrome is fast!)
        max_dim = max(img.width, img.height)
        if max_dim < 2048:
            scale_factor = 2048 / max_dim
            new_width = int(img.width * scale_factor)
            new_height = int(img.height * scale_factor)
            img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
    # Later iterations: maintain size, just enhance quality
    new_width = img.width
    new_height = img.height

    # Progressive contrast enhancement
    # Later iterations get more contrast for crisp edges
    contrast_boost = 1.3 + (iteration * 0.15)  # 1.3, 1.45, 1.6, etc. (more aggressive)
    enhancer = ImageEnhance.Contrast(img)
    img = enhancer.enhance(contrast_boost)

    # Sharpening - more aggressive in later iterations
    sharpen_radius = 2.0 + (iteration * 0.3)  # Larger radius for later iterations
    sharpen_percent = 150 + (iteration * 15)  # More aggressive sharpening
    img = img.filter(ImageFilter.UnsharpMask(radius=sharpen_radius, percent=sharpen_percent))

    # Save enhanced version
    img.save(output_path, 'PNG')

    return img.width, img.height


def trace_to_svg(png_path: str, svg_path: str, img_width: int, img_height: int, iteration: int):
    """
    Trace PNG to SVG using Inkscape.
    Progressive smoothness - higher in later iterations.
    """
    # Convert to absolute path for file:// URI (Inkscape requires absolute paths)
    abs_png_path = os.path.abspath(png_path)

    # Create temporary SVG with embedded image
    temp_svg = tempfile.NamedTemporaryFile(delete=False, suffix=".svg", mode='w')
    svg_content = f'''<?xml version="1.0" encoding="UTF-8"?>
<svg width="{img_width}" height="{img_height}" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink">
  <image id="img1" width="{img_width}" height="{img_height}" xlink:href="file://{abs_png_path}"/>
</svg>'''
    temp_svg.write(svg_content)
    temp_svg.close()

    # MINIMAL COLOR TRACING - Fast for single-color logos!
    # Use 2 scans minimum (1 scan doesn't trace properly, just fills canvas)
    scans = 2  # Minimum for proper tracing
    smooth = "true"
    stack = "true"
    remove_bg = "true"  # Remove background (transparent areas)
    speckles = 5 + (iteration * 2)  # Aggressively remove speckles in later iterations

    # Progressive smoothness - can go higher now since it's faster
    base_smoothness = 0.85 + (iteration * 0.05)  # 0.85, 0.90, 0.95
    smooth_corners = min(0.98, base_smoothness)  # Can go higher with monochrome

    # Optimization - keep minimal to preserve all curve details
    optimize = max(0.005, 0.1 - (iteration * 0.015))  # Very minimal optimization

    trace_params = f"{scans},{smooth},{stack},{remove_bg},{speckles},{smooth_corners},{optimize}"

    # Run Inkscape trace
    cmd = [
        "inkscape",
        "--batch-process",
        temp_svg.name,
        f"--actions=select-all;object-trace:{trace_params}",
        "--export-plain-svg",
        f"--export-filename={svg_path}",
    ]

    # Run Inkscape trace - NO TIMEOUT (same as web API which works)
    # This can take 2-5 minutes for large images, but it WILL complete
    console.print(f"  [dim](This may take 2-5 minutes for high quality tracing...)[/dim]")
    result = subprocess.run(cmd, capture_output=True, text=True)

    # Clean up temp file
    os.unlink(temp_svg.name)

    if result.returncode != 0:
        raise RuntimeError(f"Inkscape trace failed: {result.stderr}")

    # Clean the SVG - remove embedded image
    import re
    with open(svg_path, 'r') as f:
        svg_content = f.read()

    # Remove image elements
    svg_content = re.sub(r'<image[^>]*\/>', '', svg_content)
    svg_content = re.sub(r'<image[^>]*>.*?</image>', '', svg_content, flags=re.DOTALL)

    # Ensure proper viewBox
    if 'viewBox' not in svg_content:
        svg_content = svg_content.replace('<svg', f'<svg viewBox="0 0 {img_width} {img_height}"', 1)

    with open(svg_path, 'w') as f:
        f.write(svg_content)


def render_svg_to_png(svg_path: str, png_path: str, size: int, force_square: bool = True):
    """
    Render SVG to PNG, optionally forcing perfect square output.

    Args:
        svg_path: Path to SVG file
        png_path: Output PNG path
        size: Target size (both dimensions if force_square=True)
        force_square: If True, generates square PNG regardless of aspect ratio
    """
    if force_square:
        # Force perfect square output
        cmd = [
            "inkscape",
            svg_path,
            "--export-type=png",
            f"--export-filename={png_path}",
            f"--export-width={size}",
            f"--export-height={size}",
        ]
    else:
        # Maintain aspect ratio (original behavior)
        cmd = [
            "inkscape",
            svg_path,
            "--export-type=png",
            f"--export-filename={png_path}",
            f"--export-height={size}",
        ]

    result = subprocess.run(cmd, capture_output=True, text=True)

    if result.returncode != 0:
        raise RuntimeError(f"SVG render failed: {result.stderr}")


def iterative_refine(input_png: str, output_dir: str, base_name: str, iterations: int = 3):
    """
    Main iterative refinement pipeline.

    Args:
        input_png: Path to input PNG file
        output_dir: Directory for output files
        base_name: Base name for output files (e.g., "VeriaLogo")
        iterations: Number of refinement iterations (default: 3)

    Returns:
        Path to final refined SVG
    """
    os.makedirs(output_dir, exist_ok=True)

    current_png = input_png
    final_svg = None

    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        BarColumn(),
        TaskProgressColumn(),
        console=console
    ) as progress:

        main_task = progress.add_task(f"[cyan]Refining {base_name}...", total=iterations)

        for i in range(1, iterations + 1):
            progress.update(main_task, description=f"[cyan]Iteration {i}/{iterations}")

            # Paths for this iteration
            enhanced_png = os.path.join(output_dir, f"{base_name}_iter{i}_enhanced.png")
            iter_svg = os.path.join(output_dir, f"{base_name}_iter{i}.svg")
            rendered_png = os.path.join(output_dir, f"{base_name}_iter{i}_rendered.png")

            # Step 1: Enhance current PNG
            console.print(f"  [yellow]→[/yellow] Enhancing image (iteration {i})...")
            img_width, img_height = enhance_image(current_png, enhanced_png, i)

            # Step 2: Trace to SVG
            console.print(f"  [yellow]→[/yellow] Tracing to vector...")
            trace_to_svg(enhanced_png, iter_svg, img_width, img_height, i)
            final_svg = iter_svg

            # Step 3: Render back to PNG for next iteration (if not last)
            if i < iterations:
                console.print(f"  [yellow]→[/yellow] Rendering for next iteration...")
                # Use height as max dimension (logo is portrait orientation)
                render_svg_to_png(iter_svg, rendered_png, img_height)
                current_png = rendered_png

            progress.update(main_task, advance=1)

    console.print(f"\n[green]✓[/green] Iterative refinement complete: {final_svg}")
    return final_svg


def generate_all_sizes(final_svg: str, output_dir: str, base_name: str, sizes: list = None):
    """
    Generate optimized SVG and PNG versions at all required sizes.

    Args:
        final_svg: Path to the refined SVG
        output_dir: Output directory
        base_name: Base name for outputs
        sizes: List of sizes to generate (default: [16, 32, 40, 64, 128, 256, 512])
    """
    if sizes is None:
        sizes = [16, 32, 40, 64, 128, 256, 512]

    # Read the final SVG
    with open(final_svg, 'r') as f:
        svg_content = f.read()

    console.print(f"\n[cyan]Generating {len(sizes)} sizes...[/cyan]")

    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        BarColumn(),
        TaskProgressColumn(),
        console=console
    ) as progress:

        task = progress.add_task("[cyan]Generating sizes...", total=len(sizes))

        for size in sizes:
            # Generate optimized SVG
            svg_output = os.path.join(output_dir, f"{base_name}_{size}.svg")
            optimized_svg = optimize_for_small_size(svg_content, target_size=size)
            with open(svg_output, 'w') as f:
                f.write(optimized_svg)

            # Generate PNG (maintaining aspect ratio)
            png_output = os.path.join(output_dir, f"{base_name}_{size}.png")
            render_svg_to_png(svg_output, png_output, size)

            console.print(f"  [green]✓[/green] {size}px: {svg_output}, {png_output}")
            progress.update(task, advance=1)

    console.print(f"\n[green]✓[/green] All sizes generated in: {output_dir}")


def main():
    """CLI entry point"""
    import argparse

    parser = argparse.ArgumentParser(
        description="Iterative SVG refinement for pixel-perfect logos",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Refine logo with 3 iterations and generate all sizes
  %(prog)s input.png --output ./output --name MyLogo

  # Use 5 iterations for maximum quality
  %(prog)s input.png --output ./output --name MyLogo --iterations 5

  # Generate specific sizes only
  %(prog)s input.png --output ./output --name MyLogo --sizes 16 32 64 128
        """
    )

    parser.add_argument("input", help="Input PNG file")
    parser.add_argument("--output", "-o", default="./output", help="Output directory (default: ./output)")
    parser.add_argument("--name", "-n", default="Logo", help="Base name for output files (default: Logo)")
    parser.add_argument("--iterations", "-i", type=int, default=3, help="Number of refinement iterations (default: 3)")
    parser.add_argument("--sizes", "-s", type=int, nargs="+", help="Sizes to generate (default: 16 32 64 128 256 512)")
    parser.add_argument("--no-sizes", action="store_true", help="Skip multi-size generation")

    args = parser.parse_args()

    # Validate input
    if not os.path.exists(args.input):
        console.print(f"[red]Error: Input file not found: {args.input}[/red]")
        return 1

    # Check Inkscape
    if not ensure_inkscape():
        return 1

    console.print(f"\n[bold cyan]Iterative SVG Refinement[/bold cyan]")
    console.print(f"Input:      {args.input}")
    console.print(f"Output:     {args.output}")
    console.print(f"Name:       {args.name}")
    console.print(f"Iterations: {args.iterations}\n")

    # Run iterative refinement
    try:
        final_svg = iterative_refine(
            input_png=args.input,
            output_dir=args.output,
            base_name=args.name,
            iterations=args.iterations
        )

        # Generate all sizes
        if not args.no_sizes:
            generate_all_sizes(
                final_svg=final_svg,
                output_dir=args.output,
                base_name=args.name,
                sizes=args.sizes
            )

        console.print(f"\n[bold green]✓ Complete![/bold green]\n")
        return 0

    except Exception as e:
        console.print(f"\n[red]Error: {e}[/red]")
        import traceback
        traceback.print_exc()
        return 1


if __name__ == "__main__":
    sys.exit(main())