inference.py

"""
Restormer ONNX Inference
=========================
Standalone inference script for Restormer document denoising using ONNX Runtime.
No PyTorch required — only onnxruntime, numpy, opencv-python, and Pillow.

Usage:
    python inference.py --input noisy_doc.png --output clean_doc.png
    python inference.py --input noisy_doc.png --model fp16 --output clean_doc.png
    python inference.py --input ./noisy_dir/ --output ./clean_dir/ --batch

Models:
    fp32  — FP32 dynamic-size model (105 MB, highest fidelity)
    fp16  — FP16 single-file model (55 MB, faster, slightly reduced precision)
"""

import argparse
import os
import sys
import time
from typing import Tuple, Optional, List

import numpy as np
import onnxruntime as ort


# ---------------------------------------------------------------------------
# Paths
# ---------------------------------------------------------------------------
MODEL_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'models')

MODEL_PATHS = {
    'fp32': os.path.join(MODEL_DIR, 'restormer_denoise_dynamic.onnx'),
    'fp16': os.path.join(MODEL_DIR, 'restormer_fp16_converted.onnx'),
}

# ---------------------------------------------------------------------------
# ONNX Session
# ---------------------------------------------------------------------------
_session_cache = {}


def get_session(model: str = 'fp32', gpu: bool = True) -> ort.InferenceSession:
    """Load (or retrieve cached) ONNX Runtime inference session."""
    if model in _session_cache:
        return _session_cache[model]

    onnx_path = MODEL_PATHS.get(model)
    if onnx_path is None:
        raise ValueError(f"Unknown model '{model}'. Choose: {list(MODEL_PATHS.keys())}")
    if not os.path.exists(onnx_path):
        raise FileNotFoundError(f"ONNX model not found: {onnx_path}")

    if gpu:
        providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
    else:
        providers = ['CPUExecutionProvider']

    try:
        session = ort.InferenceSession(onnx_path, providers=providers)
    except Exception:
        print(f"⚠️  CUDA provider unavailable, falling back to CPU", file=sys.stderr)
        session = ort.InferenceSession(onnx_path, providers=['CPUExecutionProvider'])

    actual = session.get_providers()
    print(f"✅ ONNX Runtime session loaded [{model}] — providers: {actual}", file=sys.stderr)

    _session_cache[model] = session
    return session


# ---------------------------------------------------------------------------
# Tile blending
# ---------------------------------------------------------------------------

def _blend_map(tile_h: int, tile_w: int, overlap: int,
               has_left: bool, has_right: bool,
               has_top: bool, has_bottom: bool) -> np.ndarray:
    """Feathered weight map for seamless tile blending (NumPy version)."""
    weight = np.ones((tile_h, tile_w), dtype=np.float32)

    if has_left:
        ramp = np.linspace(0, 1, overlap, dtype=np.float32)
        weight[:, :overlap] *= ramp[np.newaxis, :]

    if has_right:
        ramp = np.linspace(1, 0, overlap, dtype=np.float32)
        weight[:, tile_w - overlap:] *= ramp[np.newaxis, :]

    if has_top:
        ramp = np.linspace(0, 1, overlap, dtype=np.float32)
        weight[:overlap, :] *= ramp[:, np.newaxis]

    if has_bottom:
        ramp = np.linspace(1, 0, overlap, dtype=np.float32)
        weight[tile_h - overlap:, :] *= ramp[:, np.newaxis]

    return weight


def _tile_positions(h: int, w: int, tile_size: int = 512,
                    overlap: int = 64, multiple_of: int = 8
                    ) -> List[Tuple[int, int, int, int, bool, bool, bool, bool]]:
    """Compute tile grid covering an image of size h × w."""
    tile_size = (tile_size // multiple_of) * multiple_of
    overlap = (overlap // multiple_of) * multiple_of
    stride = tile_size - overlap

    tiles = []
    y = 0
    while y < h:
        if y + tile_size > h:
            y = max(0, h - tile_size)

        x = 0
        while x < w:
            if x + tile_size > w:
                x = max(0, w - tile_size)

            tiles.append((
                y, x, tile_size, tile_size,
                x > 0, x + tile_size < w,
                y > 0, y + tile_size < h,
            ))

            if x + tile_size >= w:
                break
            x += stride

        if y + tile_size >= h:
            break
        y += stride

    return tiles


# ---------------------------------------------------------------------------
# Core inference
# ---------------------------------------------------------------------------

def denoise(image: np.ndarray, model: str = 'fp32',
            tile_size: int = 512, overlap: int = 64,
            gpu: bool = True) -> np.ndarray:
    """Denoise an RGB image (H, W, 3) uint8 using Restormer ONNX.

    Parameters
    ----------
    image : np.ndarray
        Input image in RGB format, shape (H, W, 3), dtype uint8.
    model : str
        'fp32' or 'fp16'.
    tile_size : int
        Tile size for large-image processing. Images ≤ tile_size are
        processed in a single pass.
    overlap : int
        Overlap between adjacent tiles for seamless blending.
    gpu : bool
        Prefer CUDAExecutionProvider when True.

    Returns
    -------
    np.ndarray
        Denoised image in RGB format, shape (H, W, 3), dtype uint8.
    """
    session = get_session(model, gpu=gpu)

    img_multiple_of = 8
    h, w = image.shape[:2]

    # Preprocess: [0, 255] uint8 → [0, 1] float32 NCHW
    input_ = image.astype(np.float32) / 255.0
    input_ = input_.transpose(2, 0, 1)[np.newaxis, ...]  # [1, 3, H, W]

    # Pad to multiple of 8
    H = ((h + img_multiple_of - 1) // img_multiple_of) * img_multiple_of
    W = ((w + img_multiple_of - 1) // img_multiple_of) * img_multiple_of
    pad_h, pad_w = H - h, W - w
    input_padded = np.pad(input_, ((0, 0), (0, 0), (0, pad_h), (0, pad_w)),
                          mode='reflect')

    # --- Inference ---
    if H <= tile_size and W <= tile_size:
        restored = session.run(['output'], {'input': input_padded})[0]
        restored = np.clip(restored, 0, 1)
    else:
        tiles = _tile_positions(H, W, tile_size=tile_size, overlap=overlap,
                                multiple_of=img_multiple_of)

        restored = np.zeros((1, 3, H, W), dtype=np.float32)
        weight_sum = np.zeros((1, 1, H, W), dtype=np.float32)

        for y, x, th, tw, hl, hr, ht, hb in tiles:
            tile = input_padded[:, :, y:y + th, x:x + tw]

            th2 = ((th + img_multiple_of - 1) // img_multiple_of) * img_multiple_of
            tw2 = ((tw + img_multiple_of - 1) // img_multiple_of) * img_multiple_of
            tile_padded = np.pad(tile,
                                 ((0, 0), (0, 0), (0, th2 - th), (0, tw2 - tw)),
                                 mode='reflect')

            tile_out = session.run(['output'], {'input': tile_padded})[0]
            tile_out = np.clip(tile_out, 0, 1)
            tile_out = tile_out[:, :, :th, :tw]

            blend = _blend_map(th, tw, overlap, hl, hr, ht, hb)
            blend = blend[np.newaxis, np.newaxis, ...]

            restored[:, :, y:y + th, x:x + tw] += tile_out * blend
            weight_sum[:, :, y:y + th, x:x + tw] += blend

        restored = restored / weight_sum

    # Unpad and postprocess
    restored = restored[:, :, :h, :w]
    restored = (restored[0].transpose(1, 2, 0) * 255).clip(0, 255).astype(np.uint8)

    return restored


# ---------------------------------------------------------------------------
# Utility
# ---------------------------------------------------------------------------

def denoise_file(input_path: str, output_path: str, model: str = 'fp32',
                 tile_size: int = 512, overlap: int = 64,
                 gpu: bool = True) -> None:
    """Read image from disk, denoise, and save."""
    import cv2

    img_bgr = cv2.imread(input_path)
    if img_bgr is None:
        raise FileNotFoundError(f"Cannot read image: {input_path}")

    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)

    t0 = time.perf_counter()
    result = denoise(img_rgb, model=model, tile_size=tile_size,
                     overlap=overlap, gpu=gpu)
    elapsed = (time.perf_counter() - t0) * 1000

    h, w = img_rgb.shape[:2]
    print(f"  {os.path.basename(input_path)} ({w}×{h}): {elapsed:.0f} ms [{model}]",
          file=sys.stderr)

    os.makedirs(os.path.dirname(output_path) or '.', exist_ok=True)
    result_bgr = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)
    cv2.imwrite(output_path, result_bgr)
    print(f"  Saved → {output_path}", file=sys.stderr)


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(
        description='Restormer ONNX — Document Image Denoising',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
    python inference.py -i noisy.png -o clean.png
    python inference.py -i noisy.png -o clean.png --model fp16
    python inference.py -i ./input_dir/ -o ./output_dir/ --batch
        """,
    )
    parser.add_argument('-i', '--input', required=True,
                        help='Input image path or directory (--batch mode)')
    parser.add_argument('-o', '--output', required=True,
                        help='Output image path or directory (--batch mode)')
    parser.add_argument('--model', choices=['fp32', 'fp16'], default='fp32',
                        help='Model precision: fp32 (105 MB) or fp16 (55 MB). '
                             'Default: fp32')
    parser.add_argument('--tile-size', type=int, default=512,
                        help='Tile size for large images (default: 512)')
    parser.add_argument('--overlap', type=int, default=64,
                        help='Tile overlap for blending (default: 64)')
    parser.add_argument('--cpu', action='store_true',
                        help='Force CPU inference (default: auto GPU)')
    parser.add_argument('--batch', action='store_true',
                        help='Process all images in a directory')
    args = parser.parse_args()

    # Model check
    if not os.path.exists(MODEL_PATHS[args.model]):
        print(f"❌ Model not found: {MODEL_PATHS[args.model]}", file=sys.stderr)
        print(f"   Available models: {list(MODEL_PATHS.keys())}", file=sys.stderr)
        sys.exit(1)

    use_gpu = not args.cpu

    # Batch mode
    if args.batch:
        import glob
        os.makedirs(args.output, exist_ok=True)
        exts = ('*.png', '*.jpg', '*.jpeg', '*.bmp', '*.tif', '*.tiff')
        files = []
        for ext in exts:
            files.extend(glob.glob(os.path.join(args.input, ext)))
        files = sorted(files)

        if not files:
            print(f"❌ No images found in {args.input}", file=sys.stderr)
            sys.exit(1)

        print(f"🔍 Processing {len(files)} images…", file=sys.stderr)
        for f in files:
            out_name = os.path.splitext(os.path.basename(f))[0] + '_clean.png'
            denoise_file(f, os.path.join(args.output, out_name),
                         model=args.model, tile_size=args.tile_size,
                         overlap=args.overlap, gpu=use_gpu)
    else:
        denoise_file(args.input, args.output, model=args.model,
                     tile_size=args.tile_size, overlap=args.overlap,
                     gpu=use_gpu)


if __name__ == '__main__':
    main()