rectify_with_template.py

import argparse
import sys
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parent
SRC_DIR = PROJECT_ROOT / "src"

if str(SRC_DIR) not in sys.path:
    sys.path.insert(0, str(SRC_DIR))

from warpless_docs.template_rectification import TemplateRectifier, TemplateRectifierConfig


SUPPORTED_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tif", ".tiff"}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="WarpLess Docs - template-based document rectification"
    )

    parser.add_argument(
        "--template",
        default=None,
        help="Template image path. If empty, the first image inside input/template is used.",
    )
    parser.add_argument(
        "--template-dir",
        default="input/template",
        help="Directory containing the clean flat template image.",
    )
    parser.add_argument(
        "--input",
        default=None,
        help="Single input image path. If empty, batch mode is used.",
    )
    parser.add_argument(
        "--deshadowed-dir",
        default="outputs/deshadowed",
        help="Preferred directory containing shadow-removed images.",
    )
    parser.add_argument(
        "--legacy-deshadowed-dir",
        default="outputs",
        help="Backward-compatible fallback for older outputs placed directly in outputs/.",
    )
    parser.add_argument(
        "--samples-dir",
        default="input/samples",
        help="Fallback directory containing raw sample photos.",
    )
    parser.add_argument(
        "--output-dir",
        default="outputs/template_rectified",
        help="Directory to save rectified outputs, reports, and match visualizations.",
    )
    parser.add_argument("--max-side", type=int, default=1800)
    parser.add_argument(
        "--feature-method",
        choices=["auto", "sift", "orb"],
        default="auto",
        help="Feature detector for template matching. Auto tries SIFT first, then ORB.",
    )
    parser.add_argument("--homography-only", action="store_true", help="Disable piecewise affine refinement.")
    parser.add_argument("--no-debug", action="store_true", help="Do not save match visualization images.")

    return parser.parse_args()


def is_supported_image(path: Path) -> bool:
    return path.is_file() and path.suffix.lower() in SUPPORTED_EXTENSIONS


def find_images(directory: Path, pattern: str = "*") -> list[Path]:
    if not directory.exists():
        return []
    return sorted(path for path in directory.glob(pattern) if is_supported_image(path))


def find_recursive_images(directory: Path) -> list[Path]:
    if not directory.exists():
        return []
    return sorted(path for path in directory.rglob("*") if is_supported_image(path))


def resolve_template(template_path: str | None, template_dir: Path) -> Path:
    if template_path:
        path = Path(template_path)
        if not path.exists():
            raise FileNotFoundError(f"Template image not found: {path}")
        return path

    candidates = find_images(template_dir)
    if not candidates:
        raise FileNotFoundError(
            f"No template image found in {template_dir}. "
            "Put your clean template there, for example input/template/template_page_1.png"
        )

    return candidates[0]


def resolve_inputs(args: argparse.Namespace) -> list[Path]:
    if args.input:
        path = Path(args.input)
        if not path.exists():
            raise FileNotFoundError(f"Input image not found: {path}")
        return [path]

    # Stage 2 should normally consume stage 1 outputs.
    deshadowed = find_images(Path(args.deshadowed_dir), "*_deshadowed.*")
    if deshadowed:
        print(f"Using deshadowed images from: {args.deshadowed_dir}")
        return deshadowed

    # Backward compatibility for older project versions.
    legacy = find_images(Path(args.legacy_deshadowed_dir), "*_deshadowed.*")
    if legacy:
        print(f"Using legacy deshadowed images from: {args.legacy_deshadowed_dir}")
        return legacy

    fallback = find_recursive_images(Path(args.samples_dir))
    if fallback:
        print(f"No deshadowed outputs found. Falling back to raw samples from: {args.samples_dir}")
    return fallback


def normalized_stem(input_path: Path) -> str:
    stem = input_path.stem
    if stem.endswith("_deshadowed"):
        return stem[: -len("_deshadowed")]
    return stem


def build_output_paths(input_path: Path, output_dir: Path) -> tuple[Path, Path, Path]:
    stem = normalized_stem(input_path)
    rectified_path = output_dir / f"{stem}_template_rectified.png"
    report_path = output_dir / f"{stem}_template_rectification_report.json"
    debug_path = output_dir / f"{stem}_template_matches.png"
    return rectified_path, report_path, debug_path


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

    template_path = resolve_template(args.template, Path(args.template_dir))
    input_paths = resolve_inputs(args)

    if not input_paths:
        print("No input images found.")
        print("Run shadow removal first with: python main.py")
        return

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    rectifier = TemplateRectifier(
        template_path=template_path,
        config=TemplateRectifierConfig(
            max_image_side=args.max_side,
            feature_method=args.feature_method,
            use_piecewise_warp=not args.homography_only,
            output_match_debug=not args.no_debug,
        ),
    )

    print("WarpLess Docs template rectification")
    print(f"Template  : {template_path}")
    print(f"Output dir: {output_dir}")
    print(f"Images    : {len(input_paths)}")
    print(f"Features  : {args.feature_method}")
    print(f"Mode      : {'homography only' if args.homography_only else 'homography + piecewise affine'}")
    print("-" * 78)

    for index, input_path in enumerate(input_paths, start=1):
        rectified_path, report_path, debug_path = build_output_paths(input_path, output_dir)
        print(f"[{index}/{len(input_paths)}] Processing: {input_path}")

        try:
            result = rectifier.rectify_path(
                input_path=input_path,
                output_path=rectified_path,
                output_json_path=report_path,
                output_debug_path=None if args.no_debug else debug_path,
            )
            print(
                f"  mode={result.mode} feature={result.feature_method} "
                f"good_matches={result.good_matches} inliers={result.inlier_matches} "
                f"inlier_ratio={result.inlier_ratio:.2f}"
            )
            print(f"  rectified: {rectified_path}")
            print(f"  report   : {report_path}")
            if not args.no_debug:
                print(f"  matches  : {debug_path}")
        except Exception as exc:
            print(f"  failed: {input_path}")
            print(f"  reason: {exc}")

    print("-" * 78)
    print("Done.")


if __name__ == "__main__":
    main()