diff --git a/.gitignore b/.gitignore index 880a522f..afda7ae5 100644 --- a/.gitignore +++ b/.gitignore @@ -1,18 +1,46 @@ -# Python +# Python & Virtual Environments __pycache__/ -*.pyc -*.pyo -*.pyd -.coverage +*.py[cod] +*$py.class +*.so .Python -.pytest_cache/ -.venv/ -.hypothesis env/ venv/ ENV/ env.bak/ venv.bak/ +.venv/ + +# Build & Packaging +build/ +dist/ +*.egg-info/ +*.egg +eggs/ +.eggs/ + +# Testing & Static Analysis +.pytest_cache/ +.hypothesis/ +.ruff_cache/ +.mypy_cache/ +.coverage +.coverage.* +coverage.xml +htmlcov/ + +# PyTorch, Compiler & Triton Caches +.cache/ +.triton/ + +# Evaluation & Benchmark Outputs +eval*.json +eval_output*.json +*.eval.json +benchmark_results/ +eval_results/ +tests/eval/fixtures/ +eval_fixtures/ # Project Specific ClipsForInference/* @@ -23,11 +51,6 @@ Ignored*/* !Ignored*/.gitkeep CorridorKey_remote.bat .ipynb_checkpoints/ -.DS_Store - -# IDE -.vscode/ -.idea/ # Models & Checkpoints (Large Files) *.pth @@ -49,4 +72,21 @@ gvm_core/weights/* BiRefNetModule/checkpoints/* !BiRefNetModule/checkpoints/.gitkeep -site +# Docs +site/ + +# IDEs & Editors +.vscode/ +.idea/ +*.swp +*.swo +*~ + +# OS Files +.DS_Store +.DS_Store? +._* +.Spotlight-V100 +.Trashes +ehthumbs.db +Thumbs.db diff --git a/corridorkey/__init__.py b/corridorkey/__init__.py new file mode 100644 index 00000000..0704ae28 --- /dev/null +++ b/corridorkey/__init__.py @@ -0,0 +1,3 @@ +"""CorridorKey: Neural network green screen keying for professional VFX pipelines.""" + +__version__ = "1.0.0" diff --git a/corridorkey/eval/__init__.py b/corridorkey/eval/__init__.py new file mode 100644 index 00000000..5beccf24 --- /dev/null +++ b/corridorkey/eval/__init__.py @@ -0,0 +1,34 @@ +"""Model-quality evaluation and benchmark framework for CorridorKey.""" + +from .cases import MattingEvalSample, load_dataset, load_sample_from_directory +from .fixtures import FIXTURE_GENERATORS, create_diagnostic_fixtures +from .metrics import ( + compute_alpha_mae, + compute_boundary_alpha_mae, + compute_boundary_mask, + compute_foreground_mae, + compute_recomposition_mae, +) +from .report import format_json_report, format_text_table, print_report, write_json_report +from .runner import CaseResult, EvaluationReport, evaluate_dataset, evaluate_sample + +__all__ = [ + "MattingEvalSample", + "load_dataset", + "load_sample_from_directory", + "compute_alpha_mae", + "compute_foreground_mae", + "compute_boundary_mask", + "compute_boundary_alpha_mae", + "compute_recomposition_mae", + "create_diagnostic_fixtures", + "FIXTURE_GENERATORS", + "CaseResult", + "EvaluationReport", + "evaluate_sample", + "evaluate_dataset", + "format_json_report", + "format_text_table", + "write_json_report", + "print_report", +] diff --git a/corridorkey/eval/__main__.py b/corridorkey/eval/__main__.py new file mode 100644 index 00000000..f2f79407 --- /dev/null +++ b/corridorkey/eval/__main__.py @@ -0,0 +1,132 @@ +"""Command-line entrypoint for CorridorKey model evaluation. + +Usage: + python -m corridorkey.eval --checkpoint CorridorKey_v1.0 --dataset tests/eval/fixtures --output eval.json +""" + +from __future__ import annotations + +import argparse +import os +import sys +from pathlib import Path + +from .cases import load_dataset +from .fixtures import create_diagnostic_fixtures +from .report import print_report, write_json_report +from .runner import evaluate_dataset + + +def _create_engine(checkpoint: str | None, device: str = "cpu", backend: str = "auto"): + """Instantiate inference engine based on checkpoint parameter or mock.""" + if checkpoint in ("mock", "dummy"): + # Deterministic mock engine for testing/CI without weight download + def _mock_engine(composite, hint): + # Returns composite as fg, hint as alpha + return hint, composite + + return _mock_engine + + from CorridorKeyModule.backend import create_engine + + if checkpoint and os.path.isfile(checkpoint): + from CorridorKeyModule.inference_engine import CorridorKeyEngine + + return CorridorKeyEngine(checkpoint_path=checkpoint, device=device) + + # Check if checkpoint specifies blue model + screen_color = "blue" if checkpoint and "blue" in checkpoint.lower() else "green" + return create_engine(backend=backend, device=device, screen_color=screen_color) + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser( + prog="corridorkey.eval", + description="CorridorKey Model-Quality Benchmark & Evaluation Framework", + ) + parser.add_argument( + "--checkpoint", + type=str, + default=None, + help=( + "Path to model checkpoint (.safetensors/.pth) or model name " + "('CorridorKey_v1.0', 'CorridorKeyBlue_1.0', 'mock')." + ), + ) + parser.add_argument( + "--dataset", + type=str, + default="tests/eval/fixtures", + help="Directory containing benchmark cases or a single case directory.", + ) + parser.add_argument( + "--output", + "-o", + type=str, + default=None, + help="Path to write JSON evaluation results.", + ) + parser.add_argument( + "--device", + type=str, + default="cpu", + help="Compute device ('cpu', 'cuda', 'mps').", + ) + parser.add_argument( + "--backend", + type=str, + default="auto", + choices=["auto", "torch", "mlx"], + help="Inference backend.", + ) + parser.add_argument( + "--generate-fixtures", + action="store_true", + help="Generate canonical diagnostic fixtures if dataset directory is missing or empty.", + ) + + args = parser.parse_args(argv) + + dataset_path = Path(args.dataset) + if not dataset_path.exists() or not any(dataset_path.iterdir()): + if args.generate_fixtures or "fixtures" in str(dataset_path): + print(f"Dataset path {dataset_path} empty or missing. Generating diagnostic fixtures...") + create_diagnostic_fixtures(dataset_path) + else: + print(f"Error: Dataset directory {dataset_path} does not exist.", file=sys.stderr) + return 1 + + try: + samples = load_dataset(dataset_path) + except Exception as e: + print(f"Error loading dataset from {dataset_path}: {e}", file=sys.stderr) + return 1 + + model_label = args.checkpoint or "CorridorKey_v1.0" + print(f"Loaded {len(samples)} evaluation samples from {dataset_path}") + print(f"Initializing engine ({model_label}) on {args.device}...") + + try: + engine = _create_engine(args.checkpoint, device=args.device, backend=args.backend) + except Exception as e: + print(f"Failed to initialize engine for {model_label}: {e}", file=sys.stderr) + return 1 + + print("Running evaluation suite...") + report = evaluate_dataset(engine, samples, model_name=model_label) + + # Print summary table + print_report(report) + + # Write JSON report if requested + if args.output: + out_path = Path(args.output) + out_path.parent.mkdir(parents=True, exist_ok=True) + write_json_report(report, str(out_path)) + print(f"Wrote structured evaluation receipt to {out_path}") + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/corridorkey/eval/cases.py b/corridorkey/eval/cases.py new file mode 100644 index 00000000..e05902af --- /dev/null +++ b/corridorkey/eval/cases.py @@ -0,0 +1,199 @@ +"""Dataset cases and sample contracts for matting evaluation. + +Defines the canonical MattingEvalSample and functions to load diagnostic +challenge sets from disk. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path + +import cv2 +import numpy as np +import torch + + +@dataclass(frozen=True) +class MattingEvalSample: + """Canonical evaluation sample contract for matting quality assessment. + + Attributes: + composite: Observed composite image tensor [3, H, W], float32 in [0, 1], sRGB. + alpha_hint: Coarse/guide alpha hint tensor [1, H, W], float32 in [0, 1]. + target_alpha: Ground truth alpha matte tensor [1, H, W], float32 in [0, 1]. + target_foreground: Ground truth unmultiplied straight foreground [3, H, W], float32 in [0, 1]. + screen_rgb: Optional known backing screen RGB tensor [3] or None. + tags: Frozenset of descriptive failure-mode / category tags. + name: Unique sample identifier or case name. + """ + + composite: torch.Tensor + alpha_hint: torch.Tensor + target_alpha: torch.Tensor + target_foreground: torch.Tensor + screen_rgb: torch.Tensor | None + tags: frozenset[str] + name: str = "" + + def __post_init__(self) -> None: + # Validate shapes and dimensions + if self.composite.ndim != 3 or self.composite.shape[0] != 3: + raise ValueError(f"composite must have shape [3, H, W], got {self.composite.shape}") + + h, w = self.composite.shape[1], self.composite.shape[2] + + if self.alpha_hint.ndim == 2: + object.__setattr__(self, "alpha_hint", self.alpha_hint.unsqueeze(0)) + elif self.alpha_hint.ndim != 3 or self.alpha_hint.shape[0] != 1: + raise ValueError(f"alpha_hint must have shape [1, H, W], got {self.alpha_hint.shape}") + + if self.target_alpha.ndim == 2: + object.__setattr__(self, "target_alpha", self.target_alpha.unsqueeze(0)) + elif self.target_alpha.ndim != 3 or self.target_alpha.shape[0] != 1: + raise ValueError(f"target_alpha must have shape [1, H, W], got {self.target_alpha.shape}") + + if self.target_foreground.ndim != 3 or self.target_foreground.shape[0] != 3: + raise ValueError(f"target_foreground must have shape [3, H, W], got {self.target_foreground.shape}") + + if ( + self.alpha_hint.shape[1:] != (h, w) + or self.target_alpha.shape[1:] != (h, w) + or self.target_foreground.shape[1:] != (h, w) + ): + raise ValueError( + f"Spatial dimension mismatch: composite {(h, w)}, " + f"alpha_hint {tuple(self.alpha_hint.shape[1:])}, " + f"target_alpha {tuple(self.target_alpha.shape[1:])}, " + f"target_foreground {tuple(self.target_foreground.shape[1:])}" + ) + + if self.screen_rgb is not None: + if self.screen_rgb.numel() != 3: + raise ValueError(f"screen_rgb must have 3 elements, got {self.screen_rgb.shape}") + object.__setattr__(self, "screen_rgb", self.screen_rgb.view(3).float()) + + +def _read_rgb(path: Path) -> np.ndarray: + """Read an image file as float32 RGB in range [0, 1].""" + img = cv2.imread(str(path), cv2.IMREAD_COLOR) + if img is None: + raise FileNotFoundError(f"Failed to read image at {path}") + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + return img.astype(np.float32) / 255.0 + + +def _read_grayscale(path: Path) -> np.ndarray: + """Read an image file as float32 single-channel in range [0, 1].""" + img = cv2.imread(str(path), cv2.IMREAD_GRAYSCALE) + if img is None: + raise FileNotFoundError(f"Failed to read grayscale mask at {path}") + return img.astype(np.float32) / 255.0 + + +def load_sample_from_directory(dir_path: Path | str) -> MattingEvalSample: + """Load a MattingEvalSample from a directory containing fixture assets. + + Expected directory contents: + - composite.png (or .jpg, .exr) + - alpha_hint.png (or .jpg) + - target_alpha.png (or .jpg) + - target_foreground.png (or .jpg, .exr) + - metadata.json (optional, containing 'tags' and 'screen_rgb') + + Args: + dir_path: Path to the case directory. + + Returns: + Populated MattingEvalSample. + """ + case_dir = Path(dir_path) + if not case_dir.is_dir(): + raise NotADirectoryError(f"Directory not found: {case_dir}") + + # Find composite + comp_candidates = list(case_dir.glob("composite.*")) + if not comp_candidates: + raise FileNotFoundError(f"composite image missing in {case_dir}") + comp_np = _read_rgb(comp_candidates[0]) + + # Find alpha hint + hint_candidates = list(case_dir.glob("alpha_hint.*")) + if not hint_candidates: + raise FileNotFoundError(f"alpha_hint image missing in {case_dir}") + hint_np = _read_grayscale(hint_candidates[0]) + + # Find target alpha + alpha_candidates = list(case_dir.glob("target_alpha.*")) + if not alpha_candidates: + raise FileNotFoundError(f"target_alpha image missing in {case_dir}") + alpha_np = _read_grayscale(alpha_candidates[0]) + + # Find target foreground + fg_candidates = list(case_dir.glob("target_foreground.*")) + if not fg_candidates: + raise FileNotFoundError(f"target_foreground image missing in {case_dir}") + fg_np = _read_rgb(fg_candidates[0]) + + # Load metadata if present + metadata_path = case_dir / "metadata.json" + tags: set[str] = set() + screen_rgb: torch.Tensor | None = None + + if metadata_path.is_file(): + with open(metadata_path, encoding="utf-8") as f: + meta = json.load(f) + tags = set(meta.get("tags", [])) + if "screen_rgb" in meta and meta["screen_rgb"] is not None: + screen_rgb = torch.tensor(meta["screen_rgb"], dtype=torch.float32) + + # Always include directory name as default tag if no tags defined + tags.add(case_dir.name) + + comp_tensor = torch.from_numpy(comp_np).permute(2, 0, 1).float() + hint_tensor = torch.from_numpy(hint_np).unsqueeze(0).float() + alpha_tensor = torch.from_numpy(alpha_np).unsqueeze(0).float() + fg_tensor = torch.from_numpy(fg_np).permute(2, 0, 1).float() + + return MattingEvalSample( + composite=comp_tensor, + alpha_hint=hint_tensor, + target_alpha=alpha_tensor, + target_foreground=fg_tensor, + screen_rgb=screen_rgb, + tags=frozenset(tags), + name=case_dir.name, + ) + + +def load_dataset(dataset_dir: Path | str) -> list[MattingEvalSample]: + """Load all evaluation cases from a dataset root directory. + + If dataset_dir itself is a single case directory (has composite.*), returns [sample]. + Otherwise, scans all subdirectories containing composite.*. + + Args: + dataset_dir: Path to directory containing case folders or a single case folder. + + Returns: + Sorted list of MattingEvalSample objects. + """ + root = Path(dataset_dir) + if not root.exists(): + raise FileNotFoundError(f"Dataset root not found: {root}") + + # Check if single case + if list(root.glob("composite.*")): + return [load_sample_from_directory(root)] + + samples: list[MattingEvalSample] = [] + # Discover child directories containing composite.* + for sub in sorted(root.iterdir()): + if sub.is_dir() and list(sub.glob("composite.*")): + samples.append(load_sample_from_directory(sub)) + + if not samples: + raise ValueError(f"No valid evaluation cases found under {root}") + + return samples diff --git a/corridorkey/eval/fixtures.py b/corridorkey/eval/fixtures.py new file mode 100644 index 00000000..d4b8ce25 --- /dev/null +++ b/corridorkey/eval/fixtures.py @@ -0,0 +1,368 @@ +"""Deterministic diagnostic fixture generation for matting benchmarks. + +Generates canonical test cases isolating specific failure classes: + - opaque_subject: Clean geometric silhouette, sharp boundary. + - soft_hair: Fine semi-transparent filaments with fractional alpha. + - motion_blur: Directional motion blur across transition boundary. + - green_foreground: Chromatic collision (green foreground on green backing). + - blue_foreground: Chromatic collision (blue foreground on blue backing). + - bad_hint_eroded: Severely eroded guide hint (undersegmentation recovery). + - bad_hint_dilated: Severely dilated guide hint with false-positive islands. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import cv2 +import numpy as np + + +def _save_png_rgb(path: Path, rgb: np.ndarray) -> None: + """Save float32 [0, 1] RGB array as 8-bit PNG.""" + bgr = cv2.cvtColor((np.clip(rgb, 0.0, 1.0) * 255.0).round().astype(np.uint8), cv2.COLOR_RGB2BGR) + cv2.imwrite(str(path), bgr) + + +def _save_png_gray(path: Path, gray: np.ndarray) -> None: + """Save float32 [0, 1] 2D array as 8-bit grayscale PNG.""" + u8 = (np.clip(gray, 0.0, 1.0) * 255.0).round().astype(np.uint8) + cv2.imwrite(str(path), u8) + + +def generate_opaque_subject(size: int = 256) -> dict[str, np.ndarray | dict]: + """Generate opaque geometric subject with clean boundary.""" + h, w = size, size + alpha = np.zeros((h, w), dtype=np.float32) + # Draw central rounded rectangle / circle + cv2.circle(alpha, (w // 2, h // 2), int(size * 0.35), 1.0, -1, lineType=cv2.LINE_AA) + + # Foreground: warm orange/tan tone + fg = np.zeros((h, w, 3), dtype=np.float32) + fg[:, :] = [0.85, 0.45, 0.20] + + # Screen: clean studio green + screen = np.array([0.0, 1.0, 0.0], dtype=np.float32) + screen_field = np.tile(screen, (h, w, 1)) + + # Composite: alpha * F + (1 - alpha) * B + comp = alpha[:, :, None] * fg + (1.0 - alpha[:, :, None]) * screen_field + + # Hint: slight 2px Gaussian blur of alpha + hint = cv2.GaussianBlur(alpha, (5, 5), 1.2) + + meta = { + "tags": ["opaque", "clean", "baseline"], + "screen_rgb": [0.0, 1.0, 0.0], + "description": "Clean opaque circular foreground on standard green backing.", + } + return { + "composite": comp, + "alpha_hint": hint, + "target_alpha": alpha, + "target_foreground": fg, + "metadata": meta, + } + + +def generate_soft_hair(size: int = 256) -> dict[str, np.ndarray | dict]: + """Generate subject with fine semi-transparent radial hair strands.""" + h, w = size, size + alpha = np.zeros((h, w), dtype=np.float32) + cx, cy = w // 2, int(h * 0.55) + r_core = int(size * 0.22) + + # Draw solid core + cv2.circle(alpha, (cx, cy), r_core, 1.0, -1, lineType=cv2.LINE_AA) + + # Draw 48 radial hair filaments with varying angles and lengths + rng = np.random.default_rng(1337) + for i in range(48): + angle = np.pi * 0.8 + (np.pi * 1.4) * (i / 48.0) + rng.normal(0, 0.03) + length = r_core + rng.uniform(int(size * 0.08), int(size * 0.25)) + x2 = int(cx + length * np.cos(angle)) + y2 = int(cy - length * np.sin(angle)) + # Draw line with fractional alpha + strand = np.zeros((h, w), dtype=np.float32) + cv2.line(strand, (cx, cy), (x2, y2), 1.0, thickness=1, lineType=cv2.LINE_AA) + strand = cv2.GaussianBlur(strand, (3, 3), 0.6) * rng.uniform(0.35, 0.75) + alpha = np.maximum(alpha, strand) + + alpha = np.clip(alpha, 0.0, 1.0) + + # Foreground: dark brown hair + fg = np.zeros((h, w, 3), dtype=np.float32) + fg[:, :] = [0.22, 0.14, 0.08] + + # Screen: studio green + screen = np.array([0.0, 1.0, 0.0], dtype=np.float32) + screen_field = np.tile(screen, (h, w, 1)) + + comp = alpha[:, :, None] * fg + (1.0 - alpha[:, :, None]) * screen_field + + # Hint: only covers the solid core (simulating SAM / coarse human segmenter) + hint = np.zeros((h, w), dtype=np.float32) + cv2.circle(hint, (cx, cy), int(r_core * 0.95), 1.0, -1, lineType=cv2.LINE_AA) + hint = cv2.GaussianBlur(hint, (9, 9), 2.5) + + meta = { + "tags": ["hair", "soft_hair", "fractional_alpha"], + "screen_rgb": [0.0, 1.0, 0.0], + "description": "Fine hair filaments extending into backing screen with coarse solid-core hint.", + } + return { + "composite": comp, + "alpha_hint": hint, + "target_alpha": alpha, + "target_foreground": fg, + "metadata": meta, + } + + +def generate_motion_blur(size: int = 256) -> dict[str, np.ndarray | dict]: + """Generate moving foreground with directional horizontal blur.""" + h, w = size, size + base_alpha = np.zeros((h, w), dtype=np.float32) + cv2.circle(base_alpha, (w // 2, h // 2), int(size * 0.22), 1.0, -1) + + # Horizontal motion blur kernel (length 25px) + k_len = 25 + kernel = np.zeros((k_len, k_len), dtype=np.float32) + kernel[k_len // 2, :] = 1.0 / k_len + blurred_alpha = cv2.filter2D(base_alpha, -1, kernel) + + # Foreground: magenta/red prop + fg = np.zeros((h, w, 3), dtype=np.float32) + fg[:, :] = [0.90, 0.20, 0.35] + + screen = np.array([0.0, 1.0, 0.0], dtype=np.float32) + screen_field = np.tile(screen, (h, w, 1)) + + comp = blurred_alpha[:, :, None] * fg + (1.0 - blurred_alpha[:, :, None]) * screen_field + + # Hint: static unblurred center + hint = cv2.GaussianBlur(base_alpha, (9, 9), 2.0) + + meta = { + "tags": ["motion_blur", "translucent", "temporal_spread"], + "screen_rgb": [0.0, 1.0, 0.0], + "description": "High-velocity horizontal motion blur across green backing screen.", + } + return { + "composite": comp, + "alpha_hint": hint, + "target_alpha": blurred_alpha, + "target_foreground": fg, + "metadata": meta, + } + + +def generate_green_foreground(size: int = 256) -> dict[str, np.ndarray | dict]: + """Generate green subject on green screen (severe chroma collision).""" + h, w = size, size + alpha = np.zeros((h, w), dtype=np.float32) + pts = np.array( + [ + [w * 0.5, h * 0.15], + [w * 0.8, h * 0.5], + [w * 0.65, h * 0.85], + [w * 0.35, h * 0.85], + [w * 0.2, h * 0.5], + ], + dtype=np.int32, + ) + cv2.fillPoly(alpha, [pts], 1.0, lineType=cv2.LINE_AA) + + # Green foreground (e.g. olive / lime shirt) + fg = np.zeros((h, w, 3), dtype=np.float32) + fg[:, :] = [0.20, 0.78, 0.25] + + # Screen: pure green + screen = np.array([0.0, 1.0, 0.0], dtype=np.float32) + screen_field = np.tile(screen, (h, w, 1)) + + comp = alpha[:, :, None] * fg + (1.0 - alpha[:, :, None]) * screen_field + hint = cv2.GaussianBlur(alpha, (7, 7), 1.5) + + meta = { + "tags": ["green_foreground", "chroma_collision", "adversarial_color"], + "screen_rgb": [0.0, 1.0, 0.0], + "description": "Green foliage/fabric subject against green backing screen (delta-E collision).", + } + return { + "composite": comp, + "alpha_hint": hint, + "target_alpha": alpha, + "target_foreground": fg, + "metadata": meta, + } + + +def generate_blue_foreground(size: int = 256) -> dict[str, np.ndarray | dict]: + """Generate blue subject on blue screen (severe chroma collision).""" + h, w = size, size + alpha = np.zeros((h, w), dtype=np.float32) + pts = np.array( + [ + [w * 0.5, h * 0.15], + [w * 0.85, h * 0.55], + [w * 0.65, h * 0.85], + [w * 0.35, h * 0.85], + [w * 0.15, h * 0.55], + ], + dtype=np.int32, + ) + cv2.fillPoly(alpha, [pts], 1.0, lineType=cv2.LINE_AA) + + # Blue foreground (navy/cobalt subject) + fg = np.zeros((h, w, 3), dtype=np.float32) + fg[:, :] = [0.15, 0.25, 0.88] + + # Screen: blue backing + screen = np.array([0.0, 0.0, 1.0], dtype=np.float32) + screen_field = np.tile(screen, (h, w, 1)) + + comp = alpha[:, :, None] * fg + (1.0 - alpha[:, :, None]) * screen_field + hint = cv2.GaussianBlur(alpha, (7, 7), 1.5) + + meta = { + "tags": ["blue_foreground", "chroma_collision", "adversarial_color"], + "screen_rgb": [0.0, 0.0, 1.0], + "description": "Cobalt blue foreground object against blue backing screen.", + } + return { + "composite": comp, + "alpha_hint": hint, + "target_alpha": alpha, + "target_foreground": fg, + "metadata": meta, + } + + +def generate_bad_hint_eroded(size: int = 256) -> dict[str, np.ndarray | dict]: + """Generate subject with severely eroded hint (14px undersegmentation).""" + h, w = size, size + alpha = np.zeros((h, w), dtype=np.float32) + # Head and torso + cv2.circle(alpha, (w // 2, int(h * 0.35)), int(size * 0.18), 1.0, -1, lineType=cv2.LINE_AA) + cv2.rectangle( + alpha, + (int(w * 0.25), int(h * 0.45)), + (int(w * 0.75), int(h * 0.90)), + 1.0, + -1, + lineType=cv2.LINE_AA, + ) + + fg = np.zeros((h, w, 3), dtype=np.float32) + fg[:, :] = [0.82, 0.55, 0.35] + + screen = np.array([0.0, 1.0, 0.0], dtype=np.float32) + screen_field = np.tile(screen, (h, w, 1)) + comp = alpha[:, :, None] * fg + (1.0 - alpha[:, :, None]) * screen_field + + # Heavily eroded hint: erode by 15px structuring element + k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (31, 31)) + hint_binary = cv2.erode(alpha, k) + hint = cv2.GaussianBlur(hint_binary, (9, 9), 3.0) + + meta = { + "tags": ["bad_hint", "bad_hint_eroded", "undersegmentation"], + "screen_rgb": [0.0, 1.0, 0.0], + "description": "Severely eroded alpha hint missing major perimeter regions.", + } + return { + "composite": comp, + "alpha_hint": hint, + "target_alpha": alpha, + "target_foreground": fg, + "metadata": meta, + } + + +def generate_bad_hint_dilated(size: int = 256) -> dict[str, np.ndarray | dict]: + """Generate subject with severely dilated hint and spurious false-positive islands.""" + h, w = size, size + alpha = np.zeros((h, w), dtype=np.float32) + cv2.circle(alpha, (w // 2, int(h * 0.35)), int(size * 0.18), 1.0, -1, lineType=cv2.LINE_AA) + cv2.rectangle( + alpha, + (int(w * 0.25), int(h * 0.45)), + (int(w * 0.75), int(h * 0.90)), + 1.0, + -1, + lineType=cv2.LINE_AA, + ) + + fg = np.zeros((h, w, 3), dtype=np.float32) + fg[:, :] = [0.82, 0.55, 0.35] + + screen = np.array([0.0, 1.0, 0.0], dtype=np.float32) + screen_field = np.tile(screen, (h, w, 1)) + comp = alpha[:, :, None] * fg + (1.0 - alpha[:, :, None]) * screen_field + + # Dilate by 15px structuring element + k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (31, 31)) + hint_dilated = cv2.dilate(alpha, k) + + # Add false-positive islands in backing screen + cv2.circle(hint_dilated, (int(w * 0.15), int(h * 0.2)), int(size * 0.08), 1.0, -1) + cv2.circle(hint_dilated, (int(w * 0.85), int(h * 0.75)), int(size * 0.06), 1.0, -1) + hint = cv2.GaussianBlur(hint_dilated, (9, 9), 3.0) + + meta = { + "tags": ["bad_hint", "bad_hint_dilated", "oversegmentation", "false_islands"], + "screen_rgb": [0.0, 1.0, 0.0], + "description": "Severely dilated alpha hint with extraneous false-positive background islands.", + } + return { + "composite": comp, + "alpha_hint": hint, + "target_alpha": alpha, + "target_foreground": fg, + "metadata": meta, + } + + +FIXTURE_GENERATORS = { + "opaque_subject": generate_opaque_subject, + "soft_hair": generate_soft_hair, + "motion_blur": generate_motion_blur, + "green_foreground": generate_green_foreground, + "blue_foreground": generate_blue_foreground, + "bad_hint_eroded": generate_bad_hint_eroded, + "bad_hint_dilated": generate_bad_hint_dilated, +} + + +def create_diagnostic_fixtures(target_root: Path | str, size: int = 256) -> list[Path]: + """Generate all 7 canonical benchmark fixtures into target_root. + + Args: + target_root: Destination directory for fixture folders. + size: Spatial square dimension for fixtures (default 256). + + Returns: + List of paths to created fixture directories. + """ + root = Path(target_root) + root.mkdir(parents=True, exist_ok=True) + created_paths: list[Path] = [] + + for name, gen_fn in FIXTURE_GENERATORS.items(): + case_dir = root / name + case_dir.mkdir(parents=True, exist_ok=True) + + data = gen_fn(size=size) + _save_png_rgb(case_dir / "composite.png", data["composite"]) + _save_png_gray(case_dir / "alpha_hint.png", data["alpha_hint"]) + _save_png_gray(case_dir / "target_alpha.png", data["target_alpha"]) + _save_png_rgb(case_dir / "target_foreground.png", data["target_foreground"]) + + with open(case_dir / "metadata.json", "w", encoding="utf-8") as f: + json.dump(data["metadata"], f, indent=2) + + created_paths.append(case_dir) + + return created_paths diff --git a/corridorkey/eval/metrics.py b/corridorkey/eval/metrics.py new file mode 100644 index 00000000..081e2be2 --- /dev/null +++ b/corridorkey/eval/metrics.py @@ -0,0 +1,229 @@ +"""Core evaluation metrics for matting quality assessment. + +Provides four foundational metrics: + - alpha_mae: Mean Absolute Error across all pixels for predicted alpha. + - foreground_mae: Mean Absolute Error for recovered unmultiplied straight foreground RGB. + - boundary_alpha_mae: Mean Absolute Error evaluated strictly within the transition boundary + (fractional alpha region plus a narrow morphological band), preventing easy background + pixels from masking fine-edge degradation. + - recomposition_mae: Mean Absolute Error between the observed composite and the recomposed + prediction C_hat = alpha_hat * F_hat + (1 - alpha_hat) * Screen. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + + +def _ensure_tensor(x: torch.Tensor | float | int) -> torch.Tensor: + if isinstance(x, torch.Tensor): + return x + return torch.tensor(x, dtype=torch.float32) + + +def compute_alpha_mae( + pred_alpha: torch.Tensor, + target_alpha: torch.Tensor, +) -> float: + """Mean absolute error for alpha matte. + + Args: + pred_alpha: Predicted alpha tensor, shape [..., 1, H, W] or [..., H, W], in [0, 1]. + target_alpha: Ground truth alpha tensor, matching spatial dimensions, in [0, 1]. + + Returns: + Scalar MAE as float. + """ + if pred_alpha.shape != target_alpha.shape: + if pred_alpha.ndim == target_alpha.ndim + 1 and pred_alpha.shape[-3] == 1: + pred_alpha = pred_alpha.squeeze(-3) + elif target_alpha.ndim == pred_alpha.ndim + 1 and target_alpha.shape[-3] == 1: + target_alpha = target_alpha.squeeze(-3) + else: + raise ValueError(f"Shape mismatch: pred_alpha {pred_alpha.shape} vs target_alpha {target_alpha.shape}") + + return float(torch.mean(torch.abs(pred_alpha.float() - target_alpha.float())).item()) + + +def compute_foreground_mae( + pred_fg: torch.Tensor, + target_fg: torch.Tensor, + target_alpha: torch.Tensor | None = None, + alpha_threshold: float = 0.01, +) -> float: + """Mean absolute error for straight foreground color. + + When target_alpha is provided, evaluation is restricted to pixels where + target_alpha > alpha_threshold to focus on visible foreground regions. + + Args: + pred_fg: Predicted foreground RGB [..., 3, H, W], in [0, 1]. + target_fg: Ground truth foreground RGB [..., 3, H, W], in [0, 1]. + target_alpha: Optional ground truth alpha mask [..., 1, H, W] or [..., H, W]. + alpha_threshold: Threshold below which foreground is ignored if mask is supplied. + + Returns: + Scalar MAE as float. + """ + diff = torch.abs(pred_fg.float() - target_fg.float()) + + if target_alpha is not None: + mask = target_alpha.float() > alpha_threshold + if mask.ndim == diff.ndim - 1: + mask = mask.unsqueeze(-3) + mask = mask.expand_as(diff) + if mask.sum() == 0: + return 0.0 + return float((diff * mask.float()).sum().item() / mask.sum().item()) + + return float(torch.mean(diff).item()) + + +def compute_boundary_mask( + target_alpha: torch.Tensor, + kernel_size: int = 5, + fractional_low: float = 0.01, + fractional_high: float = 0.99, +) -> torch.Tensor: + """Derive binary evaluation mask isolating difficult boundary/edge pixels. + + The boundary region includes: + 1. All fractional alpha pixels (fractional_low < alpha < fractional_high). + 2. A morphological transition band around the solid foreground boundary + (dilated binary mask minus eroded binary mask). + + Args: + target_alpha: Ground truth alpha tensor [..., H, W] or [..., 1, H, W]. + kernel_size: Size of morphological dilation/erosion structuring element (must be odd). + fractional_low: Lower bound for fractional alpha inclusion. + fractional_high: Upper bound for fractional alpha inclusion. + + Returns: + Boolean tensor of shape [..., 1, H, W] indicating boundary pixels. + """ + orig_ndim = target_alpha.ndim + if orig_ndim == 2: + alpha_4d = target_alpha.unsqueeze(0).unsqueeze(0) + elif orig_ndim == 3: + if target_alpha.shape[0] == 1: + alpha_4d = target_alpha.unsqueeze(0) + else: + alpha_4d = target_alpha.unsqueeze(1) + elif orig_ndim == 4: + alpha_4d = target_alpha + else: + raise ValueError(f"Unsupported target_alpha ndim: {orig_ndim}") + + alpha_4d = alpha_4d.float() + + # 1. Fractional alpha region + fractional = (alpha_4d > fractional_low) & (alpha_4d < fractional_high) + + # 2. Morphological boundary band + binary = (alpha_4d >= 0.5).float() + padding = kernel_size // 2 + dilated = F.max_pool2d(binary, kernel_size=kernel_size, stride=1, padding=padding) + eroded = -F.max_pool2d(-binary, kernel_size=kernel_size, stride=1, padding=padding) + morph_band = (dilated - eroded) > 0.5 + + boundary = fractional | morph_band + + if orig_ndim == 2: + return boundary.squeeze(0).squeeze(0) + elif orig_ndim == 3 and target_alpha.shape[0] != 1: + return boundary.squeeze(1) + elif orig_ndim == 3: + return boundary.squeeze(0) + return boundary + + +def compute_boundary_alpha_mae( + pred_alpha: torch.Tensor, + target_alpha: torch.Tensor, + kernel_size: int = 5, + fractional_low: float = 0.01, + fractional_high: float = 0.99, +) -> float: + """Mean absolute error of alpha restricted exclusively to the boundary region. + + Prevents models from masking poor edge/hair delineation behind large, trivial + background/foreground areas. + + Args: + pred_alpha: Predicted alpha [..., H, W] or [..., 1, H, W]. + target_alpha: Ground truth alpha [..., H, W] or [..., 1, H, W]. + kernel_size: Width of morphological edge band. + fractional_low: Fractional alpha lower threshold. + fractional_high: Fractional alpha upper threshold. + + Returns: + Scalar boundary MAE as float. + """ + if pred_alpha.ndim != target_alpha.ndim: + if pred_alpha.ndim == target_alpha.ndim + 1 and pred_alpha.shape[-3] == 1: + pred_alpha = pred_alpha.squeeze(-3) + elif target_alpha.ndim == pred_alpha.ndim + 1 and target_alpha.shape[-3] == 1: + target_alpha = target_alpha.squeeze(-3) + else: + raise ValueError(f"Shape mismatch: {pred_alpha.shape} vs {target_alpha.shape}") + + boundary_mask = compute_boundary_mask( + target_alpha, + kernel_size=kernel_size, + fractional_low=fractional_low, + fractional_high=fractional_high, + ) + + diff = torch.abs(pred_alpha.float() - target_alpha.float()) + boundary_diff = diff[boundary_mask] + + if boundary_diff.numel() == 0: + return float(torch.mean(diff).item()) + + return float(torch.mean(boundary_diff).item()) + + +def compute_recomposition_mae( + pred_alpha: torch.Tensor, + pred_fg: torch.Tensor, + composite: torch.Tensor, + screen_rgb: torch.Tensor | None = None, +) -> float: + """Mean absolute error between observed composite and recomposed prediction. + + Physical recomposition formula: + C_hat = alpha_hat * F_hat + (1 - alpha_hat) * Screen + + When screen_rgb is provided, it serves as the constant backing color B. + If screen_rgb is None, B is estimated from the composite where alpha == 0, + or directly from composite background pixels. + + Args: + pred_alpha: Predicted alpha [..., 1, H, W] or [..., H, W]. + pred_fg: Predicted straight foreground [..., 3, H, W]. + composite: Observed ground truth composite image [..., 3, H, W]. + screen_rgb: Backing screen color [3], [3, 1, 1], [3, H, W], or None. + + Returns: + Scalar recomposition MAE as float. + """ + alpha = pred_alpha.float() + if alpha.ndim == pred_fg.ndim - 1: + alpha = alpha.unsqueeze(-3) + alpha = alpha.clamp(0.0, 1.0) + fg = pred_fg.float().clamp(0.0, 1.0) + comp = composite.float() + + if screen_rgb is not None: + screen = _ensure_tensor(screen_rgb).float() + if screen.ndim == 1 and screen.shape[0] == 3: + screen = screen.view(3, 1, 1).expand_as(fg) + elif screen.ndim == 3 and screen.shape[0] == 3: + screen = screen.expand_as(fg) + else: + # If screen_rgb not provided, treat background pixels in composite as backing + screen = comp + + recomposed = alpha * fg + (1.0 - alpha) * screen + return float(torch.mean(torch.abs(recomposed - comp)).item()) diff --git a/corridorkey/eval/report.py b/corridorkey/eval/report.py new file mode 100644 index 00000000..b7e17a26 --- /dev/null +++ b/corridorkey/eval/report.py @@ -0,0 +1,116 @@ +"""Evaluation reporting and formatting utilities. + +Generates structured JSON receipts and clean terminal tables for +model benchmark results. +""" + +from __future__ import annotations + +import json +from typing import TextIO + +from .runner import EvaluationReport + + +def format_json_report(report: EvaluationReport, indent: int = 2) -> str: + """Format an EvaluationReport as a structured JSON string.""" + return json.dumps(report.to_dict(), indent=indent) + + +def write_json_report(report: EvaluationReport, out_stream: TextIO | str) -> None: + """Write an EvaluationReport to an open file stream or path string.""" + content = format_json_report(report) + if isinstance(out_stream, str): + with open(out_stream, "w", encoding="utf-8") as f: + f.write(content) + f.write("\n") + else: + out_stream.write(content) + out_stream.write("\n") + + +def format_text_table(report: EvaluationReport) -> str: + """Format an EvaluationReport as a clean ASCII benchmark table.""" + lines: list[str] = [] + lines.append("=" * 84) + header_model = f" CorridorKey Evaluation Benchmark — {report.model_name or 'Default'}" + lines.append(f"{header_model:<84}") + lines.append("=" * 84) + + # Summary section + lines.append( + f"Samples: {report.sample_count:<6} | Alpha MAE: {report.alpha_mae:.4f} | " + f"FG MAE: {report.foreground_mae:.4f} | Boundary MAE: {report.boundary_alpha_mae:.4f} | " + f"Recomp: {report.recomposition_mae:.4f}" + ) + lines.append("-" * 84) + + # Case breakdown + col_case = "Test Case" + lines.append(f"{col_case:<26} {'Alpha MAE':>12} {'FG MAE':>12} {'Boundary MAE':>14} {'Recomp MAE':>14}") + lines.append("-" * 84) + + for name, c in sorted(report.cases.items()): + lines.append( + f"{name:<26} {c.alpha_mae:>12.4f} {c.foreground_mae:>12.4f} " + f"{c.boundary_alpha_mae:>14.4f} {c.recomposition_mae:>14.4f}" + ) + + # Tag breakdown if multiple tags exist + if report.tag_breakdown: + lines.append("-" * 84) + lines.append("Failure Mode / Category Breakdown:") + lines.append( + f"{'Category Tag':<26} {'Count':>6} {'Alpha MAE':>12} {'FG MAE':>10} {'Boundary':>12} {'Recomp':>12}" + ) + lines.append("-" * 84) + for tag, metrics in sorted(report.tag_breakdown.items()): + cnt = int(metrics.get("count", 0)) + lines.append( + f"{tag:<26} {cnt:>6} {metrics['alpha_mae']:>12.4f} {metrics['foreground_mae']:>10.4f} " + f"{metrics['boundary_alpha_mae']:>12.4f} {metrics['recomposition_mae']:>12.4f}" + ) + + lines.append("=" * 84) + return "\n".join(lines) + + +def print_report(report: EvaluationReport) -> None: + """Print the benchmark report to stdout, utilizing rich if available.""" + try: + from rich.console import Console + from rich.table import Table + + console = Console() + table = Table( + title=f"CorridorKey Model Quality Benchmark: [bold green]{report.model_name or 'Candidate'}[/bold green]", + header_style="bold cyan", + ) + table.add_column("Test Case", style="bold") + table.add_column("Alpha MAE", justify="right") + table.add_column("Foreground MAE", justify="right") + table.add_column("Boundary Alpha MAE", justify="right") + table.add_column("Recomposition MAE", justify="right") + + for name, c in sorted(report.cases.items()): + table.add_row( + name, + f"{c.alpha_mae:.4f}", + f"{c.foreground_mae:.4f}", + f"{c.boundary_alpha_mae:.4f}", + f"{c.recomposition_mae:.4f}", + ) + + # Global row + table.add_section() + table.add_row( + "[bold yellow]GLOBAL MEAN[/bold yellow]", + f"[bold]{report.alpha_mae:.4f}[/bold]", + f"[bold]{report.foreground_mae:.4f}[/bold]", + f"[bold]{report.boundary_alpha_mae:.4f}[/bold]", + f"[bold]{report.recomposition_mae:.4f}[/bold]", + ) + console.print(table) + + except ImportError: + print(format_text_table(report)) diff --git a/corridorkey/eval/runner.py b/corridorkey/eval/runner.py new file mode 100644 index 00000000..d9d3965e --- /dev/null +++ b/corridorkey/eval/runner.py @@ -0,0 +1,247 @@ +"""Evaluation runner executing matting inference and aggregating metrics. + +Runs models against MattingEvalSample instances, computes core metrics, +and compiles overall, per-case, and per-tag evaluation summaries. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Sequence + +import numpy as np +import torch + +from .cases import MattingEvalSample +from .metrics import ( + compute_alpha_mae, + compute_boundary_alpha_mae, + compute_foreground_mae, + compute_recomposition_mae, +) + + +@dataclass +class CaseResult: + """Evaluation metrics for a single test case.""" + + name: str + tags: list[str] + alpha_mae: float + foreground_mae: float + boundary_alpha_mae: float + recomposition_mae: float + + def to_dict(self) -> dict[str, Any]: + return { + "alpha_mae": round(self.alpha_mae, 6), + "foreground_mae": round(self.foreground_mae, 6), + "boundary_alpha_mae": round(self.boundary_alpha_mae, 6), + "recomposition_mae": round(self.recomposition_mae, 6), + "tags": sorted(self.tags), + } + + +@dataclass +class EvaluationReport: + """Aggregated evaluation results over an entire benchmark suite.""" + + alpha_mae: float + foreground_mae: float + boundary_alpha_mae: float + recomposition_mae: float + cases: dict[str, CaseResult] = field(default_factory=dict) + tag_breakdown: dict[str, dict[str, float]] = field(default_factory=dict) + model_name: str = "" + sample_count: int = 0 + + def to_dict(self) -> dict[str, Any]: + """Serialize report to dictionary conforming to benchmark JSON schema.""" + return { + "model": self.model_name, + "sample_count": self.sample_count, + "alpha_mae": round(self.alpha_mae, 6), + "foreground_mae": round(self.foreground_mae, 6), + "boundary_alpha_mae": round(self.boundary_alpha_mae, 6), + "recomposition_mae": round(self.recomposition_mae, 6), + "cases": {name: res.to_dict() for name, res in self.cases.items()}, + "tags": {tag: {k: round(v, 6) for k, v in metrics.items()} for tag, metrics in self.tag_breakdown.items()}, + } + + +def predict_sample( + engine: Any, + sample: MattingEvalSample, +) -> tuple[torch.Tensor, torch.Tensor]: + """Execute model prediction for a sample and extract (pred_alpha, pred_fg). + + Supports: + 1. CorridorKeyEngine instance (calls process_frame). + 2. PyTorch nn.Module (forward pass with 4-channel input [C, H]). + 3. Callable (composite, hint) -> (pred_alpha, pred_fg). + + Returns: + tuple of (pred_alpha [1, H, W], pred_fg [3, H, W]) float32 tensors in [0, 1]. + """ + # 1. CorridorKeyEngine instance + if hasattr(engine, "process_frame"): + comp_np = sample.composite.permute(1, 2, 0).cpu().numpy().astype(np.float32) + hint_np = sample.alpha_hint.squeeze(0).cpu().numpy().astype(np.float32) + + screen_channel = 1 + if sample.screen_rgb is not None and sample.screen_rgb[2] > sample.screen_rgb[1]: + screen_channel = 2 + + out = engine.process_frame( + comp_np, + hint_np, + screen_channel=screen_channel, + auto_despeckle=False, + generate_comp=False, + post_process_on_gpu=False, + ) + if isinstance(out, list): + out = out[0] + + alpha_np = out["alpha"] + if alpha_np.ndim == 2: + alpha_np = alpha_np[np.newaxis, :, :] + elif alpha_np.ndim == 3 and alpha_np.shape[-1] == 1: + alpha_np = np.transpose(alpha_np, (2, 0, 1)) + + fg_np = out["fg"] + if fg_np.ndim == 3 and fg_np.shape[-1] == 3: + fg_np = np.transpose(fg_np, (2, 0, 1)) + + pred_alpha = torch.from_numpy(alpha_np).float() + pred_fg = torch.from_numpy(fg_np).float() + return pred_alpha, pred_fg + + # 2. PyTorch nn.Module + if isinstance(engine, torch.nn.Module): + engine.eval() + device = next(engine.parameters()).device + inp = torch.cat([sample.composite, sample.alpha_hint], dim=0).unsqueeze(0).to(device) + with torch.no_grad(): + out = engine(inp) + + if isinstance(out, dict): + pred_alpha = out["alpha"].detach().cpu() + pred_fg = out["fg"].detach().cpu() + elif isinstance(out, (tuple, list)): + pred_alpha, pred_fg = out[0].detach().cpu(), out[1].detach().cpu() + else: + raise TypeError(f"Unexpected module output type: {type(out)}") + + if pred_alpha.ndim == 4: + pred_alpha = pred_alpha.squeeze(0) + if pred_fg.ndim == 4: + pred_fg = pred_fg.squeeze(0) + return pred_alpha, pred_fg + + # 3. Arbitrary callable + if callable(engine): + res = engine(sample.composite, sample.alpha_hint) + if isinstance(res, dict): + pred_alpha, pred_fg = res["alpha"], res["fg"] + else: + pred_alpha, pred_fg = res[0], res[1] + + if not isinstance(pred_alpha, torch.Tensor): + pred_alpha = torch.from_numpy(np.asarray(pred_alpha)) + if not isinstance(pred_fg, torch.Tensor): + pred_fg = torch.from_numpy(np.asarray(pred_fg)) + + if pred_alpha.ndim == 2: + pred_alpha = pred_alpha.unsqueeze(0) + elif pred_alpha.ndim == 3 and pred_alpha.shape[-1] == 1: + pred_alpha = pred_alpha.permute(2, 0, 1) + + if pred_fg.ndim == 3 and pred_fg.shape[-1] == 3: + pred_fg = pred_fg.permute(2, 0, 1) + + return pred_alpha.float(), pred_fg.float() + + raise TypeError(f"Unsupported engine type: {type(engine)}") + + +def evaluate_sample(engine: Any, sample: MattingEvalSample) -> CaseResult: + """Evaluate a single sample and return individual metrics.""" + pred_alpha, pred_fg = predict_sample(engine, sample) + + # Ensure predictions are clamped to valid physical range [0, 1] + pred_alpha = pred_alpha.clamp(0.0, 1.0) + pred_fg = pred_fg.clamp(0.0, 1.0) + + alpha_mae = compute_alpha_mae(pred_alpha, sample.target_alpha) + fg_mae = compute_foreground_mae(pred_fg, sample.target_foreground, sample.target_alpha) + boundary_mae = compute_boundary_alpha_mae(pred_alpha, sample.target_alpha) + recomp_mae = compute_recomposition_mae(pred_alpha, pred_fg, sample.composite, sample.screen_rgb) + + return CaseResult( + name=sample.name or "unnamed_sample", + tags=list(sample.tags), + alpha_mae=alpha_mae, + foreground_mae=fg_mae, + boundary_alpha_mae=boundary_mae, + recomposition_mae=recomp_mae, + ) + + +def evaluate_dataset( + engine: Any, + samples: Sequence[MattingEvalSample], + model_name: str = "CorridorKey", +) -> EvaluationReport: + """Run evaluation across all provided benchmark samples and aggregate statistics. + + Args: + engine: Model runner or callable. + samples: Sequence of MattingEvalSample fixtures. + model_name: Identifier for model/checkpoint being evaluated. + + Returns: + EvaluationReport containing global, per-case, and per-tag metrics. + """ + if not samples: + raise ValueError("Cannot evaluate an empty sample list.") + + case_results: dict[str, CaseResult] = {} + tag_accum: dict[str, list[CaseResult]] = {} + + for sample in samples: + res = evaluate_sample(engine, sample) + case_results[res.name] = res + + for tag in sample.tags: + tag_accum.setdefault(tag, []).append(res) + + # Aggregate global means + n = len(samples) + global_alpha_mae = sum(c.alpha_mae for c in case_results.values()) / n + global_fg_mae = sum(c.foreground_mae for c in case_results.values()) / n + global_boundary_mae = sum(c.boundary_alpha_mae for c in case_results.values()) / n + global_recomp_mae = sum(c.recomposition_mae for c in case_results.values()) / n + + # Aggregate tag means + tag_breakdown: dict[str, dict[str, float]] = {} + for tag, c_list in sorted(tag_accum.items()): + k = len(c_list) + tag_breakdown[tag] = { + "count": float(k), + "alpha_mae": sum(c.alpha_mae for c in c_list) / k, + "foreground_mae": sum(c.foreground_mae for c in c_list) / k, + "boundary_alpha_mae": sum(c.boundary_alpha_mae for c in c_list) / k, + "recomposition_mae": sum(c.recomposition_mae for c in c_list) / k, + } + + return EvaluationReport( + model_name=model_name, + sample_count=n, + alpha_mae=global_alpha_mae, + foreground_mae=global_fg_mae, + boundary_alpha_mae=global_boundary_mae, + recomposition_mae=global_recomp_mae, + cases=case_results, + tag_breakdown=tag_breakdown, + ) diff --git a/pyproject.toml b/pyproject.toml index 29aec65f..ae320241 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,13 +67,14 @@ docs = ["zensical>=0.0.24"] [project.scripts] corridorkey = "corridorkey_cli:main" +corridorkey-eval = "corridorkey.eval.__main__:main" [build-system] requires = ["hatchling"] build-backend = "hatchling.build" [tool.hatch.build.targets.wheel] -packages = ["CorridorKeyModule", "gvm_core", "VideoMaMaInferenceModule"] +packages = ["CorridorKeyModule", "gvm_core", "VideoMaMaInferenceModule", "corridorkey", "backend"] [tool.ruff] target-version = "py311" @@ -100,7 +101,7 @@ addopts = "--tb=short" branch = true # "." covers root-level modules (clip_manager.py, corridorkey_cli.py once PR #8 merges). # Third-party and non-library code is excluded via omit. -source = ["CorridorKeyModule", "."] +source = ["CorridorKeyModule", "corridorkey", "."] omit = [ "gvm_core/*", "VideoMaMaInferenceModule/*", diff --git a/tests/eval/test_eval_cases.py b/tests/eval/test_eval_cases.py new file mode 100644 index 00000000..e700da64 --- /dev/null +++ b/tests/eval/test_eval_cases.py @@ -0,0 +1,123 @@ +"""Tests for dataset sample contracts and fixture loading.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest +import torch + +from corridorkey.eval.cases import ( + MattingEvalSample, + load_dataset, + load_sample_from_directory, +) +from corridorkey.eval.fixtures import ( + _save_png_gray, + _save_png_rgb, + create_diagnostic_fixtures, + generate_opaque_subject, +) + + +def test_matting_eval_sample_valid(): + comp = torch.rand(3, 64, 64) + hint = torch.rand(1, 64, 64) + alpha = torch.rand(1, 64, 64) + fg = torch.rand(3, 64, 64) + screen = torch.tensor([0.0, 1.0, 0.0]) + + sample = MattingEvalSample( + composite=comp, + alpha_hint=hint, + target_alpha=alpha, + target_foreground=fg, + screen_rgb=screen, + tags=frozenset({"test", "hair"}), + name="test_case", + ) + assert sample.name == "test_case" + assert "hair" in sample.tags + assert sample.composite.shape == (3, 64, 64) + assert sample.target_alpha.shape == (1, 64, 64) + + +def test_matting_eval_sample_2d_hint_squeezing(): + comp = torch.rand(3, 32, 32) + hint_2d = torch.rand(32, 32) + alpha_2d = torch.rand(32, 32) + fg = torch.rand(3, 32, 32) + + sample = MattingEvalSample( + composite=comp, + alpha_hint=hint_2d, + target_alpha=alpha_2d, + target_foreground=fg, + screen_rgb=None, + tags=frozenset(), + ) + assert sample.alpha_hint.shape == (1, 32, 32) + assert sample.target_alpha.shape == (1, 32, 32) + + +def test_matting_eval_sample_dimension_mismatch(): + comp = torch.rand(3, 32, 32) + hint = torch.rand(1, 64, 64) # Mismatched spatial size + alpha = torch.rand(1, 32, 32) + fg = torch.rand(3, 32, 32) + + with pytest.raises(ValueError, match="Spatial dimension mismatch"): + MattingEvalSample( + composite=comp, + alpha_hint=hint, + target_alpha=alpha, + target_foreground=fg, + screen_rgb=None, + tags=frozenset(), + ) + + +def test_load_sample_from_directory(tmp_path: Path): + case_dir = tmp_path / "sample_case" + case_dir.mkdir() + + data = generate_opaque_subject(size=64) + _save_png_rgb(case_dir / "composite.png", data["composite"]) + _save_png_gray(case_dir / "alpha_hint.png", data["alpha_hint"]) + _save_png_gray(case_dir / "target_alpha.png", data["target_alpha"]) + _save_png_rgb(case_dir / "target_foreground.png", data["target_foreground"]) + with open(case_dir / "metadata.json", "w", encoding="utf-8") as f: + json.dump(data["metadata"], f) + + sample = load_sample_from_directory(case_dir) + assert sample.name == "sample_case" + assert "opaque" in sample.tags + assert sample.composite.shape == (3, 64, 64) + assert sample.target_alpha.shape == (1, 64, 64) + + +def test_load_dataset_canonical_fixtures(tmp_path: Path): + fixtures_dir = tmp_path / "fixtures" + create_diagnostic_fixtures(fixtures_dir, size=64) + samples = load_dataset(fixtures_dir) + + assert len(samples) == 7 + names = {s.name for s in samples} + expected_names = { + "opaque_subject", + "soft_hair", + "motion_blur", + "green_foreground", + "blue_foreground", + "bad_hint_eroded", + "bad_hint_dilated", + } + assert names == expected_names + + for s in samples: + assert s.composite.shape == (3, 64, 64) + assert s.alpha_hint.shape == (1, 64, 64) + assert s.target_alpha.shape == (1, 64, 64) + assert s.target_foreground.shape == (3, 64, 64) + assert len(s.tags) > 0 diff --git a/tests/eval/test_eval_metrics.py b/tests/eval/test_eval_metrics.py new file mode 100644 index 00000000..04149b39 --- /dev/null +++ b/tests/eval/test_eval_metrics.py @@ -0,0 +1,134 @@ +"""Tests for matting evaluation metrics.""" + +from __future__ import annotations + +import pytest +import torch + +from corridorkey.eval.metrics import ( + compute_alpha_mae, + compute_boundary_alpha_mae, + compute_boundary_mask, + compute_foreground_mae, + compute_recomposition_mae, +) + + +def test_alpha_mae_identical(): + a = torch.rand(1, 64, 64) + assert compute_alpha_mae(a, a) == pytest.approx(0.0) + + +def test_alpha_mae_known_difference(): + a = torch.zeros(1, 32, 32) + b = torch.full((1, 32, 32), 0.25) + assert compute_alpha_mae(a, b) == pytest.approx(0.25) + + +def test_alpha_mae_dimension_squeezing(): + a_3d = torch.zeros(1, 16, 16) + b_2d = torch.full((16, 16), 0.5) + assert compute_alpha_mae(a_3d, b_2d) == pytest.approx(0.5) + assert compute_alpha_mae(b_2d, a_3d) == pytest.approx(0.5) + + +def test_alpha_mae_shape_mismatch_raises(): + a = torch.zeros(1, 16, 16) + b = torch.zeros(1, 32, 32) + with pytest.raises(ValueError, match="Shape mismatch"): + compute_alpha_mae(a, b) + + +def test_foreground_mae_identical(): + fg = torch.rand(3, 32, 32) + assert compute_foreground_mae(fg, fg) == pytest.approx(0.0) + + +def test_foreground_mae_with_alpha_mask(): + fg_pred = torch.ones(3, 10, 10) + fg_target = torch.zeros(3, 10, 10) + + # Only first column has foreground alpha > threshold + alpha = torch.zeros(1, 10, 10) + alpha[0, :, 0] = 1.0 + + # If evaluated across whole frame, error would be 1.0. + # If evaluated strictly where alpha > 0.01, it is also 1.0 on that column. + assert compute_foreground_mae(fg_pred, fg_target, alpha) == pytest.approx(1.0) + + # Now make prediction match target on the masked column, but differ on background + fg_pred[0, :, 0] = 0.0 + fg_pred[1, :, 0] = 0.0 + fg_pred[2, :, 0] = 0.0 + # Masked region has zero error, background has error 1.0. With mask, result must be 0.0! + assert compute_foreground_mae(fg_pred, fg_target, alpha) == pytest.approx(0.0) + + +def test_foreground_mae_empty_mask(): + fg1 = torch.rand(3, 10, 10) + fg2 = torch.rand(3, 10, 10) + alpha_zero = torch.zeros(1, 10, 10) + # Mask sum is 0 -> should return 0.0 cleanly without zero division + assert compute_foreground_mae(fg1, fg2, alpha_zero) == 0.0 + + +def test_compute_boundary_mask_fractional(): + # Test that fractional alpha pixels (0.01 < alpha < 0.99) are marked as boundary + alpha = torch.zeros(1, 20, 20) + alpha[0, 10, 10] = 0.5 # Fractional pixel + mask = compute_boundary_mask(alpha, kernel_size=3) + assert mask[0, 10, 10].item() is True + assert mask[0, 0, 0].item() is False + + +def test_compute_boundary_mask_step_edge(): + # Pure binary step: left half 0, right half 1 + alpha = torch.zeros(1, 20, 20) + alpha[:, :, 10:] = 1.0 + + mask = compute_boundary_mask(alpha, kernel_size=5) + # The morphological band around column 10 should be True + assert mask[0, 5, 10].item() is True + assert mask[0, 5, 9].item() is True + # Pixels far from transition should be False + assert mask[0, 5, 0].item() is False + assert mask[0, 5, 19].item() is False + + +def test_boundary_alpha_mae_isolates_edges(): + # Ground truth: 100x100 image, 10x10 square in center + target_alpha = torch.zeros(1, 100, 100) + target_alpha[0, 45:55, 45:55] = 1.0 + + # Prediction 1: perfect everywhere except 1px error right on the edge + pred_edge_err = target_alpha.clone() + pred_edge_err[0, 45, 45] = 0.0 + + # Overall MAE is tiny because of 10000 pixels + global_mae = compute_alpha_mae(pred_edge_err, target_alpha) + boundary_mae = compute_boundary_alpha_mae(pred_edge_err, target_alpha, kernel_size=3) + + assert boundary_mae > global_mae + + +def test_recomposition_mae_perfect_recomposition(): + # C = alpha * F + (1 - alpha) * B + alpha = torch.full((1, 32, 32), 0.6) + fg = torch.full((3, 32, 32), 0.8) + screen = torch.tensor([0.0, 1.0, 0.0]) + screen_field = screen.view(3, 1, 1).expand(3, 32, 32) + comp = 0.6 * fg + 0.4 * screen_field + + # Perfect prediction should yield 0.0 recomposition error + mae = compute_recomposition_mae(alpha, fg, comp, screen_rgb=screen) + assert mae == pytest.approx(0.0, abs=1e-6) + + +def test_recomposition_mae_wrong_foreground(): + alpha = torch.ones(1, 16, 16) + fg_truth = torch.ones(3, 16, 16) + fg_pred = torch.zeros(3, 16, 16) + comp = fg_truth + + mae = compute_recomposition_mae(alpha, fg_pred, comp) + assert mae == pytest.approx(1.0) diff --git a/tests/eval/test_eval_runner.py b/tests/eval/test_eval_runner.py new file mode 100644 index 00000000..87d1241a --- /dev/null +++ b/tests/eval/test_eval_runner.py @@ -0,0 +1,160 @@ +"""Tests for evaluation runner, reporting, and CLI entrypoint.""" + +from __future__ import annotations + +import json +from pathlib import Path +from unittest.mock import MagicMock + +import numpy as np +import pytest +import torch + +from corridorkey.eval.__main__ import main +from corridorkey.eval.cases import MattingEvalSample +from corridorkey.eval.report import format_json_report, format_text_table, write_json_report +from corridorkey.eval.runner import evaluate_dataset, evaluate_sample, predict_sample + + +def _create_dummy_sample(name: str = "dummy", tags: tuple[str, ...] = ("test",)) -> MattingEvalSample: + h, w = 32, 32 + comp = torch.zeros(3, h, w) + hint = torch.zeros(1, h, w) + alpha = torch.zeros(1, h, w) + alpha[0, 10:20, 10:20] = 1.0 + fg = torch.full((3, h, w), 0.8) + screen = torch.tensor([0.0, 1.0, 0.0]) + return MattingEvalSample( + composite=comp, + alpha_hint=hint, + target_alpha=alpha, + target_foreground=fg, + screen_rgb=screen, + tags=frozenset(tags), + name=name, + ) + + +def test_predict_sample_callable(): + sample = _create_dummy_sample() + + def dummy_fn(comp, hint): + return hint, comp + + pred_alpha, pred_fg = predict_sample(dummy_fn, sample) + assert pred_alpha.shape == (1, 32, 32) + assert pred_fg.shape == (3, 32, 32) + + +def test_predict_sample_engine_mock(): + sample = _create_dummy_sample() + mock_engine = MagicMock() + mock_engine.process_frame.return_value = { + "alpha": np.zeros((32, 32, 1), dtype=np.float32), + "fg": np.zeros((32, 32, 3), dtype=np.float32), + } + + pred_alpha, pred_fg = predict_sample(mock_engine, sample) + assert pred_alpha.shape == (1, 32, 32) + assert pred_fg.shape == (3, 32, 32) + assert mock_engine.process_frame.called + + +def test_predict_sample_torch_module(): + sample = _create_dummy_sample() + + class TinyModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(4, 4, 1) + + def forward(self, x): + out = self.conv(x) + return {"alpha": out[:, :1], "fg": out[:, 1:]} + + mod = TinyModule() + pred_alpha, pred_fg = predict_sample(mod, sample) + assert pred_alpha.shape == (1, 32, 32) + assert pred_fg.shape == (3, 32, 32) + + +def test_evaluate_sample(): + sample = _create_dummy_sample(name="sample_a") + + # Oracle engine that returns ground truth + def oracle_engine(comp, hint): + return sample.target_alpha, sample.target_foreground + + res = evaluate_sample(oracle_engine, sample) + assert res.name == "sample_a" + assert res.alpha_mae == pytest.approx(0.0) + assert res.foreground_mae == pytest.approx(0.0) + assert res.boundary_alpha_mae == pytest.approx(0.0) + + +def test_evaluate_dataset_aggregation(): + sample1 = _create_dummy_sample(name="case_1", tags=("hair", "fine")) + sample2 = _create_dummy_sample(name="case_2", tags=("hair", "motion")) + + def stub_engine(comp, hint): + # Off-by-one everywhere + return sample1.target_alpha + 0.1, sample1.target_foreground + 0.2 + + report = evaluate_dataset(stub_engine, [sample1, sample2], model_name="StubModel") + assert report.sample_count == 2 + assert report.model_name == "StubModel" + assert "case_1" in report.cases + assert "case_2" in report.cases + assert "hair" in report.tag_breakdown + assert report.tag_breakdown["hair"]["count"] == 2.0 + + +def test_report_formatting(tmp_path: Path): + sample = _create_dummy_sample(name="case_1") + report = evaluate_dataset(lambda c, h: (h, c), [sample], model_name="TestModel") + + # Test text formatting + text_tbl = format_text_table(report) + assert "TestModel" in text_tbl + assert "case_1" in text_tbl + + # Test JSON formatting + json_str = format_json_report(report) + data = json.loads(json_str) + assert data["model"] == "TestModel" + assert "alpha_mae" in data + assert "cases" in data + + # Test file writing + out_file = tmp_path / "report.json" + write_json_report(report, str(out_file)) + assert out_file.is_file() + + +def test_cli_execution_mock(tmp_path: Path): + dataset_dir = tmp_path / "fixtures" + out_file = tmp_path / "cli_eval.json" + exit_code = main( + [ + "--checkpoint", + "mock", + "--dataset", + str(dataset_dir), + "--generate-fixtures", + "--output", + str(out_file), + "--device", + "cpu", + ] + ) + assert exit_code == 0 + assert out_file.is_file() + + with open(out_file, encoding="utf-8") as f: + data = json.load(f) + + assert data["sample_count"] == 7 + assert "alpha_mae" in data + assert "cases" in data + assert "soft_hair" in data["cases"] + assert "opaque_subject" in data["cases"]