chromainvert 0.1.0__py3-none-any.whl

chromainvert/__init__.py

@@ -0,0 +1,3 @@
+from chromainvert.cube import CubeLut, identity_lattice, read_cube, write_cube
+
+__all__ = ["CubeLut", "identity_lattice", "read_cube", "write_cube"]

chromainvert/__main__.py

@@ -0,0 +1,4 @@
+from chromainvert.cli import main
+
+if __name__ == "__main__":
+    raise SystemExit(main())

chromainvert/benchmark.py

@@ -0,0 +1,206 @@
+import time
+
+import numpy as np
+import torch
+from numpy.typing import NDArray
+
+from chromainvert.color import delta_e_2000, linear_srgb_to_lab
+from chromainvert.cube import CubeLut
+from chromainvert.lut_module import DifferentiableLUT3D
+from chromainvert.optimize import OptimizeConfig, optimize_lut
+
+
+def hsv_to_rgb(hsv: NDArray[np.floating]) -> NDArray[np.float64]:
+    hsv = np.asarray(hsv, dtype=np.float64)
+    h = (hsv[..., 0] * 6.0) % 6.0
+    s = hsv[..., 1]
+    v = hsv[..., 2]
+    i = np.floor(h).astype(np.int64)
+    f = h - np.floor(h)
+    p = v * (1.0 - s)
+    q = v * (1.0 - f * s)
+    t = v * (1.0 - (1.0 - f) * s)
+    out = np.empty(hsv.shape)
+    table = [
+        (v, t, p),
+        (q, v, p),
+        (p, v, t),
+        (p, q, v),
+        (t, p, v),
+        (v, p, q),
+    ]
+    for k, (r, g, b) in enumerate(table):
+        m = i == k
+        out[m, 0], out[m, 1], out[m, 2] = r[m], g[m], b[m]
+    return out
+
+
+def make_chart(width: int = 1536, height: int = 512, seed: int = 7) -> NDArray[np.float64]:
+    rng = np.random.default_rng(seed)
+    yy, xx = np.mgrid[0:height, 0:width].astype(np.float64)
+    chart = np.zeros((height, width, 3))
+    chart[..., 0] = 0.05 + 0.9 * xx / width
+    chart[..., 1] = 0.05 + 0.9 * yy / height
+    chart[..., 2] = 0.05 + 0.9 * ((xx + yy) / (width + height))
+    rows, cols = 4, 8
+    ph, pw = height * 3 // 4 // rows, width // cols
+    for r in range(rows):
+        for c in range(cols):
+            hue = (r * cols + c) / (rows * cols)
+            patch = hsv_to_rgb(np.array([[[hue, 0.85, 0.45 + 0.5 * ((r + c) % 2)]]]))
+            chart[r * ph : (r + 1) * ph, c * pw : (c + 1) * pw] = patch
+    cy, cx, ry, rx = height // 3, width // 6, height // 8, width // 20
+    ellipse = ((yy - cy) / ry) ** 2 + ((xx - cx) / rx) ** 2 <= 1.0
+    chart[ellipse] = (0.78, 0.57, 0.47)
+    ramp = np.linspace(0.0, 1.0, width).reshape(1, width, 1)
+    chart[height * 7 // 8 :, :, :] = ramp
+    chart = chart + rng.normal(0.0, 0.005, chart.shape)
+    return np.clip(chart, 0.0, 1.0)
+
+
+def commercial_style_lattice(size: int = 33) -> NDArray[np.float64]:
+    coords = np.linspace(0.0, 1.0, size)
+    grid = np.stack(np.meshgrid(coords, coords, coords, indexing="ij"), axis=-1)[..., ::-1]
+    rgb = grid * np.array([1.08, 1.0, 0.92]) + np.array([0.02, 0.01, 0.03])
+    luma = rgb @ np.array([0.2126, 0.7152, 0.0722])
+    scale = (luma + 0.5 * luma * (1.0 - luma) * (luma - 0.5)) / np.clip(luma, 1e-6, None)
+    rgb = rgb * scale[..., None]
+    luma = np.clip(rgb @ np.array([0.2126, 0.7152, 0.0722]), 0.0, 1.0)
+    rgb = luma[..., None] + 1.25 * (rgb - luma[..., None])
+    shadow_w = (1.0 - luma) ** 2
+    highlight_w = luma**2
+    rgb = rgb + shadow_w[..., None] * np.array([-0.03, 0.01, 0.04])
+    rgb = rgb + highlight_w[..., None] * np.array([0.04, 0.01, -0.03])
+    return np.ascontiguousarray(np.clip(rgb, 0.0, 1.0), dtype=np.float64)
+
+
+def apply_lut(image: NDArray[np.floating], lattice: NDArray[np.floating]) -> NDArray[np.float32]:
+    module = DifferentiableLUT3D(int(lattice.shape[0]), lattice=np.asarray(lattice, dtype=np.float32))
+    with torch.no_grad():
+        return module(torch.from_numpy(np.asarray(image, dtype=np.float32))).numpy()
+
+
+def mean_de00(a: NDArray[np.floating], b: NDArray[np.floating], stride: int = 16) -> float:
+    a = np.asarray(a).reshape(-1, 3)[::stride]
+    b = np.asarray(b).reshape(-1, 3)[::stride]
+    return float(delta_e_2000(linear_srgb_to_lab(a), linear_srgb_to_lab(b)).mean())
+
+
+def mean_de00_full(a: NDArray[np.floating], b: NDArray[np.floating]) -> float:
+    return mean_de00(a, b, stride=1)
+
+
+def random_rgb_probe(n: int = 10000, seed: int = 99) -> NDArray[np.float32]:
+    if n <= 0:
+        raise ValueError("probe_count must be positive")
+    rng = np.random.default_rng(seed)
+    return np.ascontiguousarray(rng.uniform(0.0, 1.0, size=(n, 3)), dtype=np.float32)
+
+
+def style_probe_metrics(
+    probe: NDArray[np.floating],
+    truth_lattice: NDArray[np.floating],
+    recovered_lattice: NDArray[np.floating],
+) -> dict[str, float]:
+    probe_metrics = heldout_probe_metrics(probe, truth_lattice, recovered_lattice)
+    return {
+        "style_probe_de00": probe_metrics["recovered_probe_de00"],
+        "style_baseline_de00": probe_metrics["baseline_probe_de00"],
+        "style_improvement": probe_metrics["improvement_probe"],
+    }
+
+
+def is_probe_gate_pass(baseline_probe: float, recovered_probe: float, improvement_probe: float, threshold: float = 0.4) -> bool:
+    # Absolute fallback requires strict improvement: recovered must beat baseline and sit below 3.0.
+    return bool(improvement_probe >= threshold or (recovered_probe < 3.0 and recovered_probe < baseline_probe))
+
+
+def heldout_probe_metrics(
+    probe: NDArray[np.floating],
+    truth_lattice: NDArray[np.floating],
+    recovered_lattice: NDArray[np.floating],
+) -> dict[str, float]:
+    probe_arr = np.asarray(probe, dtype=np.float32).reshape(-1, 3)
+    truth_probe = apply_lut(probe_arr[:, None, :], truth_lattice).reshape(-1, 3)
+    recovered_probe = apply_lut(probe_arr[:, None, :], recovered_lattice).reshape(-1, 3)
+    baseline = mean_de00_full(truth_probe, probe_arr)
+    recovered_de = mean_de00_full(truth_probe, recovered_probe)
+    improvement = (baseline - recovered_de) / baseline if baseline > 0 else 0.0
+    return {
+        "baseline_probe_de00": baseline,
+        "recovered_probe_de00": recovered_de,
+        "improvement_probe": improvement,
+    }
+
+
+def run_benchmark(
+    coarse_size: int = 17,
+    fine_size: int = 33,
+    iterations: int = 400,
+    fine_iterations: int = 30,
+    seed: int = 123,
+    device: str = "auto",
+    pixel_budget: int = 65536,
+    deterministic: bool = False,
+    probe_count: int = 10000,
+    probe_seed: int | None = None,
+) -> dict:
+    if probe_count <= 0:
+        raise ValueError("probe_count must be positive")
+    started = time.monotonic()
+    chart = make_chart(seed=seed)
+    other = make_chart(seed=seed + 1)
+    heldout = make_chart(seed=seed + 2)
+    source = chart
+    truth = CubeLut(size=33, lattice=commercial_style_lattice(33).astype(np.float32))
+    reference = apply_lut(other, truth.lattice)
+    if probe_seed is None:
+        probe_seed = seed + 100
+    config = OptimizeConfig(
+        coarse_size=coarse_size,
+        fine_size=fine_size,
+        iterations=iterations,
+        fine_iterations=fine_iterations,
+        seed=seed,
+        device=device,
+        pixel_budget=pixel_budget,
+        deterministic=deterministic,
+    )
+    lut, report = optimize_lut([source], [reference], config)
+    graded_heldout = apply_lut(heldout, lut.lattice)
+    truth_heldout = apply_lut(heldout, truth.lattice)
+    baseline = mean_de00(truth_heldout, heldout)
+    recovered = mean_de00(truth_heldout, graded_heldout)
+    rmse = float(np.sqrt(np.mean((truth_heldout - graded_heldout) ** 2)))
+    probe = random_rgb_probe(probe_count, probe_seed)
+    probe_metrics = heldout_probe_metrics(probe, truth.lattice, lut.lattice)
+    baseline_probe = probe_metrics["baseline_probe_de00"]
+    recovered_probe = probe_metrics["recovered_probe_de00"]
+    improvement_probe = probe_metrics["improvement_probe"]
+    improvement_heldout = (baseline - recovered) / baseline if baseline > 0 else 0.0
+    chart_gate_pass = bool(improvement_heldout >= 0.4)
+    probe_gate_pass = is_probe_gate_pass(baseline_probe, recovered_probe, improvement_probe)
+    return {
+        "seed": seed,
+        "coarse_size": coarse_size,
+        "fine_size": fine_size,
+        "iterations": iterations,
+        "fine_iterations": fine_iterations,
+        "baseline_heldout_de00": baseline,
+        "recovered_heldout_de00": recovered,
+        "improvement_heldout": improvement_heldout,
+        "heldout_psnr": 20.0 * np.log10(1.0 / rmse) if rmse > 0 else float("inf"),
+        "seconds": time.monotonic() - started,
+        "optimizer_report": report,
+        "baseline_probe_de00": baseline_probe,
+        "recovered_probe_de00": recovered_probe,
+        "improvement_probe": improvement_probe,
+        "probe_count": probe_count,
+        "probe_seed": probe_seed,
+        "heldout_gate_pass": probe_gate_pass,
+        "probe_gate_pass": probe_gate_pass,
+        "chart_gate_pass": chart_gate_pass,
+        "benchmark_gate_pass": bool(chart_gate_pass and probe_gate_pass),
+        "heldout_gate_threshold": 0.4,
+        "heldout_gate_absolute": 3.0,
+    }

chromainvert/cli.py

@@ -0,0 +1,256 @@
+import argparse
+import json
+import sys
+import time
+from pathlib import Path
+
+import imageio.v2 as imageio
+import numpy as np
+import torch
+
+from chromainvert.color import decode_input, encode_input, premiere_apple_log_to_working
+from chromainvert.cube import CubeLut, identity_lattice, write_cube
+from chromainvert.lut_module import DifferentiableLUT3D
+from chromainvert.lut_predictor import LutPredictor, grade_features, lattice_to_cube
+from chromainvert.optimize import OptimizeConfig, optimize_lut
+
+_ALLOWED_PROFILES = {"rec709", "rec2020", "apple_log", "slog3", "slog3_sgamut3cine", "slog3_sgamut3.cine", "slog3_sgamut3_cine"}
+_LOG_PROFILES = {"apple_log", "slog3", "slog3_sgamut3cine", "slog3_sgamut3.cine", "slog3_sgamut3_cine"}
+
+
+def default_reference_profile(profile: str) -> str:
+    return "rec709" if profile in _LOG_PROFILES else profile
+
+
+def _canonical_profile(profile: str) -> str:
+    return "slog3" if profile in _LOG_PROFILES and profile != "apple_log" else profile
+
+
+def bake_export_lut(lut: CubeLut, input_profile: str, reference_profile: str, export_target: str = "log") -> CubeLut:
+    if input_profile in _LOG_PROFILES and _canonical_profile(reference_profile) == _canonical_profile(input_profile):
+        return lut
+    codes = identity_lattice(lut.size).astype(np.float64)
+    working = premiere_apple_log_to_working(codes) if export_target == "premiere" else decode_input(codes, input_profile)
+    module = DifferentiableLUT3D(lut.size, lattice=lut.lattice)
+    with torch.no_grad():
+        graded = module(torch.from_numpy(np.ascontiguousarray(working, dtype=np.float32))).numpy()
+    display = encode_input(graded.astype(np.float64), reference_profile)
+    lattice = np.ascontiguousarray(np.clip(display, 0.0, 1.0), dtype=np.float32)
+    return CubeLut(size=lut.size, lattice=lattice, title=lut.title)
+
+
+def resolve(base: Path, p: str) -> Path:
+    path = Path(p)
+    return path if path.is_absolute() else base / path
+
+
+def load_frames(paths: list[str], profile: str, workdir: Path) -> list[np.ndarray]:
+    frames: list[np.ndarray] = []
+    for p in paths:
+        raw = imageio.imread(resolve(workdir, p))
+        if raw.ndim == 2:
+            raw = np.stack((raw, raw, raw), axis=-1)
+        if raw.ndim != 3 or raw.shape[-1] not in (3, 4):
+            raise ValueError(f"expected an RGB or RGBA image, got shape {raw.shape}")
+        if np.issubdtype(raw.dtype, np.integer):
+            scale = float(np.iinfo(raw.dtype).max)
+        elif np.issubdtype(raw.dtype, np.floating):
+            scale = 1.0
+        else:
+            raise ValueError(f"unsupported image dtype: {raw.dtype}")
+        image = raw[..., :3].astype(np.float64) / scale
+        frames.append(decode_input(image, profile))
+    return frames
+
+
+def blend_toward_identity(lut: CubeLut, strength: float) -> CubeLut:
+    if not 0.0 <= strength <= 1.0:
+        raise ValueError("strength must be in [0, 1]")
+    lattice = (strength * lut.lattice + (1.0 - strength) * identity_lattice(lut.size)).astype(np.float32)
+    return CubeLut(size=lut.size, lattice=lattice, title=lut.title)
+
+
+def cmd_run(args: argparse.Namespace) -> int:
+    job_path = Path(args.job)
+    workdir = job_path.parent
+    job = json.loads(job_path.read_text(encoding="utf-8"))
+    contract_version = job.get("contract_version")
+    if type(contract_version) is not int or contract_version not in (1, 2):
+        raise ValueError("unsupported contract_version, expected 1 or 2")
+    for key in ("reference_frames", "output_cube", "report_json"):
+        if key not in job:
+            raise ValueError(f"job is missing required key: {key}")
+    reference_only = bool(job.get("reference_only", False))
+    sources = job.get("source_frames", [])
+    if reference_only:
+        if sources:
+            raise ValueError("source_frames must be omitted when reference_only is true")
+        if "predict_weights" not in job:
+            raise ValueError("job is missing required key: predict_weights")
+    elif "source_frames" not in job:
+        raise ValueError("job is missing required key: source_frames")
+    config = OptimizeConfig(
+        coarse_size=int(job.get("coarse_size", 17)),
+        fine_size=int(job.get("lut_size", 33)),
+        iterations=int(job.get("iterations", 400)),
+        fine_iterations=int(job.get("fine_iterations", 30)),
+        seed=int(job.get("seed", 42)),
+        device=str(job.get("device", "auto")),
+        deterministic=bool(job.get("deterministic", False)),
+        weight_tv=float(job.get("weight_tv", 0.1)),
+        weight_shadow=float(job.get("weight_shadow", 0.25)),
+        weight_paired=float(job.get("weight_paired", 0.0)),
+    )
+    profile = str(job.get("input_profile", "rec709"))
+    if profile not in _ALLOWED_PROFILES:
+        raise ValueError(f"unsupported input_profile {profile!r}, expected one of {sorted(_ALLOWED_PROFILES)}")
+    if contract_version == 1 and profile not in ("rec709", "rec2020", "apple_log"):
+        raise ValueError(f"input_profile {profile!r} requires contract_version 2")
+    reference_profile = str(job.get("reference_profile", default_reference_profile(profile)))
+    if reference_profile not in _ALLOWED_PROFILES:
+        raise ValueError(f"unsupported reference_profile {reference_profile!r}, expected one of {sorted(_ALLOWED_PROFILES)}")
+    if reference_profile in _LOG_PROFILES and _canonical_profile(reference_profile) != _canonical_profile(profile):
+        raise ValueError(f"reference_profile {reference_profile!r} is a log profile; graded references must use a display profile such as 'rec709'")
+    export_target = str(job.get("export_target", "log"))
+    if export_target not in ("log", "premiere"):
+        raise ValueError(f"unsupported export_target {export_target!r}, expected 'log' or 'premiere'")
+    if export_target == "premiere" and (profile != "apple_log" or reference_profile in _LOG_PROFILES):
+        raise ValueError("export_target 'premiere' requires input_profile 'apple_log' and a display reference_profile")
+    strength = float(job.get("strength", 1.0))
+    if not 0.0 <= strength <= 1.0:
+        raise ValueError("strength must be in [0, 1]")
+    source = load_frames(job["source_frames"], profile, workdir) if not reference_only else []
+    reference = load_frames(job["reference_frames"], reference_profile, workdir)
+    started = time.monotonic()
+
+    def progress(i: int, total: int, loss: float, elapsed: float) -> None:
+        print(json.dumps({"event": "progress", "iter": i, "iters": total, "loss": loss, "elapsed_s": elapsed}), flush=True)
+
+    if reference_only:
+        lut = predict_for_frames(reference, str(job["predict_weights"]), profile)
+        report = {"mode": "single_image", "predict_weights": str(job["predict_weights"])}
+    else:
+        lut, report = optimize_lut(source, reference, config, progress=progress)
+        report["mode"] = "paired"
+    lut = blend_toward_identity(lut, strength)
+    lut = bake_export_lut(lut, profile, reference_profile, export_target)
+    lut.title = "ChromaInvert"
+    cube_path = resolve(workdir, job["output_cube"])
+    write_cube(cube_path, lut)
+    report["contract_version"] = contract_version
+    report["input_profile"] = profile
+    report["reference_profile"] = reference_profile
+    report["weight_tv"] = config.weight_tv
+    report["weight_shadow"] = config.weight_shadow
+    report["weight_paired"] = config.weight_paired
+    report["export_domain"] = "working" if profile in _LOG_PROFILES and _canonical_profile(reference_profile) == _canonical_profile(profile) else "host"
+    report["export_target"] = export_target
+    report["strength"] = strength
+    report["versions"] = {"python": sys.version.split()[0], "numpy": np.__version__, "torch": torch.__version__}
+    resolve(workdir, job["report_json"]).write_text(json.dumps(report, indent=2))
+    print(json.dumps({"event": "done", "cube": str(cube_path), "seconds": time.monotonic() - started}), flush=True)
+    return 0
+
+
+def predict_for_frames(reference_frames: list[np.ndarray], weights: str, profile: str) -> CubeLut:
+    if not reference_frames:
+        raise ValueError("predict needs at least one reference frame")
+    stacked = np.concatenate([np.asarray(f, dtype=np.float32).reshape(-1, 3) for f in reference_frames], axis=0)
+    features = torch.from_numpy(grade_features(stacked)).unsqueeze(0)
+    checkpoint = torch.load(weights, map_location="cpu")
+    if isinstance(checkpoint, dict) and "state_dict" in checkpoint:
+        model = LutPredictor(hidden=tuple(checkpoint.get("hidden", (128, 256))))
+        model.load_state_dict(checkpoint["state_dict"])
+    else:
+        model = LutPredictor()
+        model.load_state_dict(checkpoint)
+    model.eval()
+    with torch.no_grad():
+        lattice = model(features).squeeze(0).numpy()
+    return lattice_to_cube(np.ascontiguousarray(lattice, dtype=np.float32))
+
+
+def cmd_predict(args: argparse.Namespace) -> int:
+    raw = imageio.imread(args.image)
+    if raw.ndim == 2:
+        raw = np.stack((raw, raw, raw), axis=-1)
+    if raw.ndim != 3 or raw.shape[-1] not in (3, 4):
+        raise ValueError(f"expected an RGB or RGBA image, got shape {raw.shape}")
+    if np.issubdtype(raw.dtype, np.integer):
+        pixels = raw[..., :3].astype(np.float64) / float(np.iinfo(raw.dtype).max)
+    else:
+        pixels = raw[..., :3].astype(np.float64)
+    profile = str(args.profile)
+    if profile not in _ALLOWED_PROFILES:
+        raise ValueError(f"unsupported input_profile {profile!r}")
+    reference_profile = default_reference_profile(profile)
+    working = np.asarray(decode_input(pixels, reference_profile), dtype=np.float32)
+    started = time.monotonic()
+    cube = predict_for_frames([working], str(args.weights), profile)
+    cube = bake_export_lut(cube, profile, reference_profile)
+    cube_path = Path(args.out)
+    write_cube(cube_path, cube)
+    print(json.dumps({"event": "done", "cube": str(cube_path), "seconds": time.monotonic() - started}), flush=True)
+    return 0
+
+
+def cmd_benchmark(args: argparse.Namespace) -> int:
+    from chromainvert.benchmark import run_benchmark
+
+    metrics = run_benchmark(
+        coarse_size=args.coarse,
+        fine_size=args.fine,
+        iterations=args.iterations,
+        fine_iterations=args.fine_iterations,
+        seed=args.seed,
+        device=args.device,
+        pixel_budget=args.pixels,
+        deterministic=args.deterministic,
+        probe_count=args.probe_count,
+        probe_seed=args.probe_seed,
+    )
+    out = Path(args.out)
+    out.parent.mkdir(parents=True, exist_ok=True)
+    out.write_text(json.dumps(metrics, indent=2))
+    print(f"baseline held-out mean dE00: {metrics['baseline_heldout_de00']:.3f}")
+    print(f"recovered held-out mean dE00: {metrics['recovered_heldout_de00']:.3f}")
+    print(f"improvement: {100.0 * metrics['improvement_heldout']:.1f}%")
+    print(f"baseline probe dE00: {metrics['baseline_probe_de00']:.3f}")
+    print(f"recovered probe dE00: {metrics['recovered_probe_de00']:.3f}")
+    print(f"probe improvement: {100.0 * metrics['improvement_probe']:.1f}%")
+    print(f"probe gate: {'PASS' if metrics['heldout_gate_pass'] else 'FAIL'}")
+    print(f"seconds: {metrics['seconds']:.1f}")
+    return 0 if bool(metrics.get("benchmark_gate_pass", False)) else 1
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(prog="chromainvert")
+    sub = parser.add_subparsers(dest="command", required=True)
+    p_run = sub.add_parser("run", help="optimize a LUT from a job file")
+    p_run.add_argument("job")
+    p_run.set_defaults(func=cmd_run)
+    p_pred = sub.add_parser("predict", help="predict a LUT from a single graded image")
+    p_pred.add_argument("image")
+    p_pred.add_argument("--profile", default="rec709")
+    p_pred.add_argument("--weights", required=True)
+    p_pred.add_argument("--out", required=True)
+    p_pred.set_defaults(func=cmd_predict)
+    p_bench = sub.add_parser("benchmark", help="synthetic ground-truth recovery check")
+    p_bench.add_argument("--out", default="work/benchmark.json")
+    p_bench.add_argument("--coarse", type=int, default=17)
+    p_bench.add_argument("--fine", type=int, default=33)
+    p_bench.add_argument("--iterations", type=int, default=400)
+    p_bench.add_argument("--fine-iterations", type=int, default=30)
+    p_bench.add_argument("--seed", type=int, default=123)
+    p_bench.add_argument("--device", default="auto")
+    p_bench.add_argument("--pixels", type=int, default=65536)
+    p_bench.add_argument("--deterministic", action="store_true")
+    p_bench.add_argument("--probe-count", type=int, default=10000, dest="probe_count")
+    p_bench.add_argument("--probe-seed", type=int, default=None, dest="probe_seed")
+    p_bench.set_defaults(func=cmd_benchmark)
+    parsed = parser.parse_args(argv)
+    return parsed.func(parsed)
+
+
+if __name__ == "__main__":
+    sys.exit(main())