main.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.shadow_removal import DocShadowONNXRemover


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


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="WarpLess Docs - ML document shadow removal")

    parser.add_argument(
        "--input",
        default=None,
        help="Single input image path. If empty, framed outputs are processed first.",
    )
    parser.add_argument(
        "--samples-dir",
        default="outputs/framed",
        help="Preferred directory containing frame-corrected document photos.",
    )
    parser.add_argument(
        "--fallback-samples-dir",
        default="input/samples",
        help="Fallback directory containing raw document photos if outputs/framed is empty.",
    )
    parser.add_argument(
        "--output-dir",
        default="outputs/deshadowed",
        help="Directory to save shadow-removed images.",
    )
    parser.add_argument(
        "--model",
        default="models/docshadow_sd7k.onnx",
        help="Path to DocShadow ONNX model.",
    )
    parser.add_argument(
        "--size",
        nargs=2,
        type=int,
        default=[768, 768],
        metavar=("WIDTH", "HEIGHT"),
        help="Inference resize size. Try 1024 1024 for higher quality.",
    )
    parser.add_argument("--cpu", action="store_true", help="Force CPU provider.")

    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 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]

    framed = find_images(Path(args.samples_dir))
    if framed:
        print(f"Using frame-corrected images from: {args.samples_dir}")
        return framed

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

    return []


def build_output_path(input_path: Path, output_dir: Path) -> Path:
    return output_dir / f"{input_path.stem}_deshadowed.png"


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

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

    remover = DocShadowONNXRemover(
        model_path=model_path,
        input_size=(args.size[0], args.size[1]),
        prefer_gpu=not args.cpu,
    )

    input_paths = resolve_inputs(args)

    if not input_paths:
        print("No input images found.")
        print("Run frame correction first with: python scan_document_frame.py")
        return

    print("WarpLess Docs shadow removal")
    print(f"Model     : {model_path}")
    print(f"Output dir: {output_dir}")
    print(f"Providers : {remover.providers}")
    print(f"Images    : {len(input_paths)}")
    print("-" * 68)

    for index, input_path in enumerate(input_paths, start=1):
        if not input_path.exists():
            print(f"[SKIP] Input not found: {input_path}")
            continue

        output_path = build_output_path(input_path, output_dir)
        print(f"[{index}/{len(input_paths)}] Processing: {input_path}")

        try:
            saved_path = remover.remove_shadow_from_path(input_path=input_path, output_path=output_path)
            print(f"  saved: {saved_path}")
        except Exception as exc:
            print(f"  failed: {input_path}")
            print(f"  reason: {exc}")

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


if __name__ == "__main__":
    main()