scan_document_frame.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.document_frame import DocumentFrameRectifier, FrameRectifierConfig


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


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="WarpLess Docs - document frame detection and perspective correction")
    parser.add_argument("--input", default=None, help="Single input image path. If empty, all images in input/samples are processed.")
    parser.add_argument("--samples-dir", default="input/samples")
    parser.add_argument("--output-dir", default="outputs/framed")
    parser.add_argument("--debug-dir", default="outputs/frame_detection")
    parser.add_argument("--max-side", type=int, default=1800)
    parser.add_argument("--min-area", type=float, default=0.18, help="Minimum page area ratio for contour detection.")
    parser.add_argument("--no-debug", action="store_true")
    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) -> list[Path]:
    if not directory.exists():
        return []
    return sorted(path for path in directory.rglob("*") if is_supported_image(path))


def build_paths(input_path: Path, output_dir: Path, debug_dir: Path) -> tuple[Path, Path, Path]:
    output_path = output_dir / f"{input_path.stem}_framed.png"
    report_path = debug_dir / f"{input_path.stem}_frame_report.json"
    debug_path = debug_dir / f"{input_path.stem}_frame_debug.png"
    return output_path, report_path, debug_path


def main() -> None:
    args = parse_args()
    output_dir = Path(args.output_dir)
    debug_dir = Path(args.debug_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    debug_dir.mkdir(parents=True, exist_ok=True)

    if args.input:
        input_paths = [Path(args.input)]
    else:
        input_paths = find_images(Path(args.samples_dir))

    if not input_paths:
        print("No input images found.")
        return

    rectifier = DocumentFrameRectifier(
        FrameRectifierConfig(
            max_image_side=args.max_side,
            min_area_ratio=args.min_area,
        )
    )

    print("WarpLess Docs frame rectification")
    print(f"Output dir: {output_dir}")
    print(f"Debug dir : {debug_dir}")
    print(f"Images    : {len(input_paths)}")
    print("-" * 76)

    for index, input_path in enumerate(input_paths, start=1):
        output_path, report_path, debug_path = build_paths(input_path, output_dir, debug_dir)
        print(f"[{index}/{len(input_paths)}] Processing: {input_path}")
        try:
            result = rectifier.rectify_path(
                input_path=input_path,
                output_path=output_path,
                output_json_path=report_path,
                output_debug_path=None if args.no_debug else debug_path,
            )
            print(
                f"  status={result.status} method={result.method} "
                f"confidence={result.confidence:.2f} area={result.document_area_ratio:.2f}"
            )
            print(f"  framed: {output_path}")
            print(f"  report: {report_path}")
            if not args.no_debug:
                print(f"  debug : {debug_path}")
        except Exception as exc:
            print(f"  failed: {input_path}")
            print(f"  reason: {exc}")

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


if __name__ == "__main__":
    main()