batch_errormap.py

import argparse
import os
import re
from collections import defaultdict
from dataclasses import dataclass, field
from pathlib import Path
from typing import DefaultDict, Iterable

import numpy as np

try:
    from shadowmap_ops import read_exr, write_exr
except ModuleNotFoundError:
    from process_exr.shadowmap_ops import read_exr, write_exr


SCRIPT_DIR = Path(__file__).resolve().parent
FRAME_PATTERN = re.compile(r"\.(\d+)\.exr$", re.IGNORECASE)
METRIC_CHOICES = ("e", "ae", "se", "rae", "rse")


@dataclass(slots=True)
class MatchGroup:
    image_matches: list[Path] = field(default_factory=list)
    reference_matches: list[Path] = field(default_factory=list)


@dataclass(slots=True)
class ScanStats:
    total_exr_files: int = 0
    skipped_no_frame: int = 0
    matched_pairs: int = 0
    missing_image: int = 0
    missing_reference: int = 0
    duplicate_image: int = 0
    duplicate_reference: int = 0
    written_files: int = 0


def parse_args() -> argparse.Namespace:
    default_input = SCRIPT_DIR / "input"
    default_output = SCRIPT_DIR / "output" / "errormap"

    parser = argparse.ArgumentParser(
        description="按同目录同帧号匹配两类 EXR，并输出 tev 风格的 errormap",
    )
    parser.add_argument(
        "--input",
        type=Path,
        default=default_input,
        help=f"输入根目录，默认: {default_input}",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=default_output,
        help=f"输出根目录，默认: {default_output}",
    )
    parser.add_argument(
        "--image-contains",
        required=True,
        metavar="STR",
        help="待比较图像文件名必须包含的子串，例如 ours",
    )
    parser.add_argument(
        "--reference-contains",
        required=True,
        metavar="STR",
        help="参考图像文件名必须包含的子串，例如 gt；相对误差以它为分母",
    )
    parser.add_argument(
        "--metric",
        choices=METRIC_CHOICES,
        default="rae",
        help="误差算法：e, ae, se, rae, rse；默认 rae，与 tev 的 Relative Absolute Error 一致",
    )
    parser.add_argument(
        "--epsilon",
        type=float,
        default=0.01,
        help="相对误差分母中的稳定项，默认 0.01，与 tev 默认公式一致",
    )
    parser.add_argument(
        "--suffix",
        default="",
        help="追加到输出文件名中的自定义后缀，例如 exp1；为空则只使用 metric 名称",
    )
    parser.add_argument(
        "--no-recursive",
        action="store_true",
        help="不递归子目录，默认递归",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="只打印配对结果与统计，不写出 EXR",
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="打印每个配对和输出路径",
    )
    return parser.parse_args()


def iter_exr_files(root: Path, recursive: bool) -> Iterable[Path]:
    if recursive:
        for current_root, _, files in os.walk(root):
            current_path = Path(current_root)
            for name in files:
                if name.lower().endswith(".exr"):
                    yield current_path / name
        return

    for child in root.iterdir():
        if child.is_file() and child.suffix.lower() == ".exr":
            yield child


def extract_frame_number(path: Path) -> str | None:
    match = FRAME_PATTERN.search(path.name)
    if not match:
        return None
    return match.group(1)


def build_groups(
    input_root: Path,
    image_contains: str,
    reference_contains: str,
    recursive: bool,
) -> tuple[DefaultDict[tuple[str, str], MatchGroup], ScanStats]:
    groups: DefaultDict[tuple[str, str], MatchGroup] = defaultdict(MatchGroup)
    stats = ScanStats()

    image_contains_lower = image_contains.lower()
    reference_contains_lower = reference_contains.lower()

    for file_path in iter_exr_files(input_root, recursive):
        stats.total_exr_files += 1
        frame = extract_frame_number(file_path)
        if frame is None:
            stats.skipped_no_frame += 1
            continue

        relative_parent = file_path.relative_to(input_root).parent.as_posix()
        group = groups[(relative_parent, frame)]
        file_name_lower = file_path.name.lower()
        if image_contains_lower in file_name_lower:
            group.image_matches.append(file_path)
        if reference_contains_lower in file_name_lower:
            group.reference_matches.append(file_path)

    return groups, stats


def apply_metric(image: np.ndarray, reference: np.ndarray, metric: str, epsilon: float) -> np.ndarray:
    if image.shape != reference.shape:
        raise ValueError(f"shape 不一致: image={image.shape}, reference={reference.shape}")

    image = image.astype(np.float32, copy=False)
    reference = reference.astype(np.float32, copy=False)
    diff = image - reference

    if metric == "e":
        result = diff
    elif metric == "ae":
        result = np.abs(diff)
    elif metric == "se":
        result = diff * diff
    elif metric == "rae":
        result = np.abs(diff) / (reference + float(epsilon))
    elif metric == "rse":
        result = (diff * diff) / (reference * reference + float(epsilon))
    else:
        raise ValueError(f"不支持的 metric: {metric}")

    if result.ndim == 3 and result.shape[2] >= 4:
        result[:, :, 3] = 1.0

    return result.astype(np.float32, copy=False)


def replace_first_ignore_case(source: str, old: str, new: str) -> str:
    pattern = re.compile(re.escape(old), re.IGNORECASE)
    return pattern.sub(new, source, count=1)


def build_output_name(reference_path: Path, metric: str, reference_contains: str, suffix: str) -> str:
    frame = extract_frame_number(reference_path)
    if frame is None:
        raise ValueError(f"无法从参考文件提取帧号: {reference_path}")

    prefix = reference_path.name[: -(len(frame) + len(".exr") + 1)]
    suffix_parts = ["errormap", metric]
    if suffix:
        suffix_parts.append(suffix)
    replacement = "_".join(suffix_parts)

    if reference_contains and reference_contains.lower() in prefix.lower():
        output_prefix = replace_first_ignore_case(prefix, reference_contains, replacement)
    else:
        output_prefix = f"{prefix}_{replacement}"

    return f"{output_prefix}.{frame}.exr"


def print_summary(stats: ScanStats) -> None:
    print("Summary:")
    print(f"  total_exr_files={stats.total_exr_files}")
    print(f"  skipped_no_frame={stats.skipped_no_frame}")
    print(f"  matched_pairs={stats.matched_pairs}")
    print(f"  missing_image={stats.missing_image}")
    print(f"  missing_reference={stats.missing_reference}")
    print(f"  duplicate_image={stats.duplicate_image}")
    print(f"  duplicate_reference={stats.duplicate_reference}")
    print(f"  written_files={stats.written_files}")


def main() -> None:
    args = parse_args()

    input_root = args.input.resolve()
    output_root = args.output.resolve()
    recursive = not args.no_recursive

    if not input_root.exists():
        raise SystemExit(f"input path not found: {input_root}")
    if not input_root.is_dir():
        raise SystemExit(f"input path is not a directory: {input_root}")
    if args.epsilon <= 0:
        raise SystemExit("epsilon 必须大于 0")
    if args.image_contains.lower() == args.reference_contains.lower():
        raise SystemExit("image-contains 和 reference-contains 不能相同")

    groups, stats = build_groups(
        input_root=input_root,
        image_contains=args.image_contains,
        reference_contains=args.reference_contains,
        recursive=recursive,
    )

    for (relative_parent, frame), group in sorted(groups.items()):
        if not group.image_matches and not group.reference_matches:
            continue

        if len(group.image_matches) == 0:
            stats.missing_image += 1
            if args.verbose:
                print(f"[skip] missing image for frame={frame}, dir={relative_parent or '.'}")
            continue
        if len(group.reference_matches) == 0:
            stats.missing_reference += 1
            if args.verbose:
                print(f"[skip] missing reference for frame={frame}, dir={relative_parent or '.'}")
            continue
        if len(group.image_matches) > 1:
            stats.duplicate_image += 1
            if args.verbose:
                print(f"[skip] duplicate image matches for frame={frame}, dir={relative_parent or '.'}")
                for item in group.image_matches:
                    print(f"        image: {item}")
            continue
        if len(group.reference_matches) > 1:
            stats.duplicate_reference += 1
            if args.verbose:
                print(f"[skip] duplicate reference matches for frame={frame}, dir={relative_parent or '.'}")
                for item in group.reference_matches:
                    print(f"        reference: {item}")
            continue

        image_path = group.image_matches[0]
        reference_path = group.reference_matches[0]
        output_name = build_output_name(reference_path, args.metric, args.reference_contains, args.suffix)
        output_path = output_root / Path(relative_parent) / output_name if relative_parent else output_root / output_name

        stats.matched_pairs += 1
        if args.verbose or args.dry_run:
            print(f"[match] {image_path}  <->  {reference_path}")
            print(f"        output: {output_path}")

        if args.dry_run:
            continue

        image = read_exr(image_path)
        reference = read_exr(reference_path)
        errormap = apply_metric(image, reference, args.metric, args.epsilon)
        write_exr(output_path, errormap)
        stats.written_files += 1

    print_summary(stats)


if __name__ == "__main__":
    main()