fredcallagan/uvdoc-grid-onnx

1#!/usr/bin/env python32"""3UVDoc Grid-Output Document Unwarping Example4 5This script demonstrates how to use the UVDoc ONNX model with grid output6for high-resolution document unwarping.7 8The key advantage of this grid-output model over image-output models is that9the coordinate grid can be upscaled to any resolution, preserving document10quality when applied via cv2.remap().11 12Usage:13    python example.py input_image.jpg output_image.jpg14    python example.py input_image.jpg output_image.jpg --model path/to/UVDoc_grid.onnx15 16Requirements:17    pip install onnxruntime opencv-python numpy18 19Optional (for automatic model download):20    pip install huggingface_hub21"""22 23import argparse24import sys25from pathlib import Path26 27import cv228import numpy as np29 30# Model input dimensions (fixed for UVDoc architecture)31MODEL_INPUT_HEIGHT = 72032MODEL_INPUT_WIDTH = 49633 34 35def load_model(model_path: str = None):36    """37    Load the ONNX model.38 39    Args:40        model_path: Path to the ONNX model file. If None, attempts to download41                    from HuggingFace Hub.42 43    Returns:44        ONNX Runtime InferenceSession45    """46    import onnxruntime as ort47 48    if model_path is None:49        try:50            from huggingface_hub import hf_hub_download51 52            print("Downloading model from HuggingFace Hub...")53            model_path = hf_hub_download(54                repo_id="YOUR_USERNAME/uvdoc-grid-onnx",  # Update with actual repo55                filename="UVDoc_grid.onnx"56            )57            print(f"Model downloaded to: {model_path}")58        except ImportError:59            print("Error: huggingface_hub not installed. Install it or provide --model path.")60            print("  pip install huggingface_hub")61            sys.exit(1)62 63    print(f"Loading model from: {model_path}")64    session = ort.InferenceSession(65        model_path,66        providers=['CPUExecutionProvider']67    )68 69    return session70 71 72def preprocess_image(image: np.ndarray) -> np.ndarray:73    """74    Preprocess image for UVDoc model input.75 76    Args:77        image: BGR image from cv2.imread()78 79    Returns:80        Preprocessed tensor of shape (1, 3, 720, 496)81    """82    # Convert BGR to RGB83    img_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)84 85    # Resize to model input size86    resized = cv2.resize(img_rgb, (MODEL_INPUT_WIDTH, MODEL_INPUT_HEIGHT))87 88    # Normalize to [0, 1]89    normalized = resized.astype(np.float32) / 255.090 91    # Convert HWC to CHW format92    transposed = np.transpose(normalized, (2, 0, 1))93 94    # Add batch dimension95    batched = np.expand_dims(transposed, axis=0)96 97    return batched98 99 100def apply_grid_unwarping(101    image: np.ndarray,102    grid: np.ndarray,103    interpolation: int = cv2.INTER_CUBIC104) -> np.ndarray:105    """106    Apply the coordinate grid to unwarp the image.107 108    Args:109        image: Original BGR image (any resolution)110        grid: Model output grid of shape (1, 2, 45, 31)111        interpolation: OpenCV interpolation method112 113    Returns:114        Unwarped image at original resolution115    """116    h_orig, w_orig = image.shape[:2]117 118    # Remove batch dimension and transpose to (H, W, 2)119    grid_2d = np.transpose(grid[0], (1, 2, 0))  # (45, 31, 2)120 121    # Upscale grid to original image resolution122    grid_upscaled = cv2.resize(123        grid_2d,124        (w_orig, h_orig),125        interpolation=cv2.INTER_LINEAR126    )127 128    # Convert normalized coordinates [-1, 1] to pixel coordinates129    # Grid channel 0 = x (width), channel 1 = y (height)130    map_x = ((grid_upscaled[..., 0] + 1) / 2) * (w_orig - 1)131    map_y = ((grid_upscaled[..., 1] + 1) / 2) * (h_orig - 1)132 133    # Apply remapping134    unwarped = cv2.remap(135        image,136        map_x.astype(np.float32),137        map_y.astype(np.float32),138        interpolation=interpolation,139        borderMode=cv2.BORDER_REPLICATE140    )141 142    return unwarped143 144 145def unwarp_document(146    image_path: str,147    output_path: str,148    model_path: str = None149) -> None:150    """151    Main function to unwarp a document image.152 153    Args:154        image_path: Path to input warped document image155        output_path: Path to save unwarped result156        model_path: Optional path to ONNX model file157    """158    # Load image159    print(f"Loading image: {image_path}")160    image = cv2.imread(image_path)161    if image is None:162        print(f"Error: Could not load image from {image_path}")163        sys.exit(1)164 165    h, w = image.shape[:2]166    print(f"Image size: {w}x{h}")167 168    # Load model169    session = load_model(model_path)170 171    # Get input name172    input_name = session.get_inputs()[0].name173    print(f"Model input name: {input_name}")174 175    # Preprocess176    print("Preprocessing image...")177    input_tensor = preprocess_image(image)178    print(f"Input tensor shape: {input_tensor.shape}")179 180    # Run inference181    print("Running inference...")182    result = session.run(None, {input_name: input_tensor})[0]183    print(f"Output grid shape: {result.shape}")184    print(f"Output grid range: [{result.min():.4f}, {result.max():.4f}]")185 186    # Apply unwarping187    print("Applying grid-based unwarping...")188    unwarped = apply_grid_unwarping(image, result)189 190    # Save result191    print(f"Saving result to: {output_path}")192    cv2.imwrite(output_path, unwarped)193 194    print("Done!")195 196 197def main():198    parser = argparse.ArgumentParser(199        description="Unwarp document images using UVDoc grid-output ONNX model",200        formatter_class=argparse.RawDescriptionHelpFormatter,201        epilog="""202Examples:203    python example.py warped_doc.jpg unwarped_doc.jpg204    python example.py warped_doc.jpg unwarped_doc.jpg --model UVDoc_grid.onnx205        """206    )207 208    parser.add_argument(209        "input",210        help="Path to input warped document image"211    )212 213    parser.add_argument(214        "output",215        help="Path to save unwarped output image"216    )217 218    parser.add_argument(219        "--model", "-m",220        default=None,221        help="Path to UVDoc_grid.onnx model file (downloads from HuggingFace if not provided)"222    )223 224    args = parser.parse_args()225 226    # Validate input file exists227    if not Path(args.input).exists():228        print(f"Error: Input file not found: {args.input}")229        sys.exit(1)230 231    unwarp_document(args.input, args.output, args.model)232 233 234if __name__ == "__main__":235    main()236