Deploy inference-only Gradio demo

.gitattributes

@@ -33,3 +33,8 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/COCO_val2014_000000016995.jpg filter=lfs diff=lfs merge=lfs -text
+examples/COCO_val2014_000000122166.jpg filter=lfs diff=lfs merge=lfs -text
+examples/COCO_val2014_000000168337.jpg filter=lfs diff=lfs merge=lfs -text
+examples/COCO_val2014_000000446053.jpg filter=lfs diff=lfs merge=lfs -text
+examples/COCO_val2014_000000477919.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,12 +1,15 @@
 ---
 title: Image Rotation Angle Estimation
-emoji: 👀
-colorFrom: red
-colorTo: blue
+emoji: "\U0001F504"
+colorFrom: blue
+colorTo: green
 sdk: gradio
-sdk_version: 6.10.0
+sdk_version: "5.23.0"
 app_file: app.py
+python_version: "3.11"
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Upload an image to predict and correct its rotation angle. Uses the CGD (Circular Gaussian Distribution) method with a MambaOut Base backbone, achieving 2.84° MAE on COCO.
+
+[GitHub](https://github.com/maxwoe/image-rotation-angle-estimation) | [Paper](https://arxiv.org/abs/2603.25351) | [Pretrained Models](https://huggingface.co/maxwoe/image-rotation-angle-estimation)

app.py

@@ -0,0 +1,174 @@
+"""HuggingFace Spaces demo for Image Rotation Angle Estimation.
+
+Two-step interactive demo:
+1. Upload an image and click "Random Rotate" to apply a random rotation
+2. Click "Correct Orientation" to see the model predict and correct the angle
+"""
+
+import json
+
+import gradio as gr
+import torch
+from PIL import Image
+import os
+import random
+import numpy as np
+from loguru import logger
+from huggingface_hub import hf_hub_download
+
+from model_cgd import CGDAngleEstimation
+from architectures import get_default_input_size
+from rotation_utils import rotate_image_crop_max_area
+
+# HuggingFace Hub configuration
+HF_REPO_ID = os.environ.get("HF_MODEL_REPO", "maxwoe/image-rotation-angle-estimation")
+
+# Fetch config.json from model repo (this is the HF download-tracking query file)
+_config_path = hf_hub_download(repo_id=HF_REPO_ID, filename="config.json")
+with open(_config_path) as _f:
+    _config = json.load(_f)
+HF_MODELS = _config["models"]
+HF_DEFAULT_MODEL = _config["default_model"]
+
+# Global model state
+model = None
+current_model_name = None
+
+
+def get_device():
+    return "cuda:0" if torch.cuda.is_available() else "cpu"
+
+
+def load_model(name):
+    """Download and load a model from HuggingFace Hub."""
+    global model, current_model_name
+    if name == current_model_name and model is not None:
+        return gr.Info(f"Model already loaded: {name}")
+
+    if name not in HF_MODELS:
+        return gr.Warning(f"Unknown model: {name}")
+
+    info = HF_MODELS[name]
+    logger.info(f"Downloading {info['filename']} from {HF_REPO_ID}...")
+    local_path = hf_hub_download(repo_id=HF_REPO_ID, filename=info["filename"])
+
+    architecture = info["architecture"]
+    image_size = get_default_input_size(architecture)
+
+    logger.info(f"Loading model from {local_path}...")
+    new_model = CGDAngleEstimation.try_load(checkpoint_path=local_path, image_size=image_size)
+    new_model.eval()
+
+    device = get_device()
+    if device.startswith("cuda"):
+        new_model = new_model.to(device)
+
+    model = new_model
+    current_model_name = name
+    logger.info(f"Model loaded: {name} on {device}")
+    return gr.Info(f"Loaded: {name} ({device})")
+
+
+def store_original(image):
+    """Store the uploaded image as the original for rotation."""
+    if image is None:
+        return None, ""
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+    return image, ""
+
+
+def random_rotate(original):
+    """Apply a random rotation to the original uploaded image."""
+    if original is None:
+        return None, None, ""
+
+    angle = random.uniform(0, 360)
+    img_array = np.array(original)
+    rotated_array = rotate_image_crop_max_area(img_array, angle)
+    rotated = Image.fromarray(rotated_array)
+    return rotated, angle, f"{angle:.1f}°"
+
+
+def correct_orientation(image):
+    """Predict the rotation angle and correct the image."""
+    if image is None:
+        return None, "Please upload and rotate an image first."
+    if model is None:
+        return None, "Model is still loading, please wait..."
+
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+
+    predicted_angle = model.predict_angle(image)
+
+    corrected = image.rotate(-predicted_angle, expand=True, fillcolor=(255, 255, 255))
+
+    return corrected, f"Predicted rotation: {predicted_angle:.2f}°"
+
+
+# Build UI
+app = gr.Blocks(title="Image Rotation Angle Estimation")
+with app:
+    gr.HTML("<h1>Image Rotation Angle Estimation</h1>")
+    gr.Markdown(
+        "Upload an image, apply a random rotation, and see the model predict and correct the angle.\n\n"
+        "Uses the **CGD** (Circular Gaussian Distribution) method with **MambaOut Base** architecture. "
+    )
+
+    original_image_state = gr.State(value=None)
+    actual_angle_state = gr.State(value=None)
+
+    model_dropdown = gr.Dropdown(
+        choices=list(HF_MODELS.keys()),
+        value=HF_DEFAULT_MODEL,
+        label="Model",
+    )
+    model_dropdown.change(load_model, inputs=[model_dropdown])
+
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(type="pil", label="Upload Image", height=400, format="png")
+            rotate_btn = gr.Button("Random Rotate", variant="secondary", size="lg")
+        with gr.Column():
+            corrected_image = gr.Image(label="Corrected Image", height=400, interactive=False, format="png")
+            correct_btn = gr.Button("Correct Orientation", variant="primary", size="lg")
+
+    with gr.Row():
+        rotation_info = gr.Textbox(label="Applied Rotation", lines=1, interactive=False)
+        result_text = gr.Textbox(label="Predicted Rotation", lines=1, interactive=False)
+
+    gr.Examples(
+        examples=[
+            "examples/COCO_val2014_000000168337.jpg",
+            "examples/COCO_val2014_000000122166.jpg",
+            "examples/COCO_val2014_000000446053.jpg",
+            "examples/COCO_val2014_000000477919.jpg",
+            "examples/COCO_val2014_000000016995.jpg",
+        ],
+        inputs=input_image,
+        fn=store_original,
+        outputs=[original_image_state, rotation_info],
+        cache_examples=False,
+        run_on_click=True,
+    )
+
+    input_image.upload(
+        store_original,
+        inputs=[input_image],
+        outputs=[original_image_state, rotation_info],
+    )
+    rotate_btn.click(
+        random_rotate,
+        inputs=[original_image_state],
+        outputs=[input_image, actual_angle_state, rotation_info],
+    )
+    correct_btn.click(
+        correct_orientation,
+        inputs=[input_image],
+        outputs=[corrected_image, result_text],
+    )
+
+    app.load(lambda: load_model(HF_DEFAULT_MODEL))
+
+app.launch()

architectures.py

@@ -0,0 +1,25 @@
+"""Architecture configurations for inference (input sizes only)."""
+
+ARCHITECTURES = {
+    "vit_tiny_patch16_224.augreg_in21k_ft_in1k": {"input_size": 224},
+    "vit_base_patch16_224.augreg_in21k_ft_in1k": {"input_size": 224},
+    "efficientvit_b0.r224_in1k": {"input_size": 224},
+    "efficientvit_b3.r224_in1k": {"input_size": 224},
+    "convnextv2_atto.fcmae_ft_in1k": {"input_size": 224},
+    "convnextv2_base.fcmae_ft_in22k_in1k": {"input_size": 224},
+    "efficientnetv2_rw_t.ra2_in1k": {"input_size": 224},
+    "efficientnetv2_rw_m.agc_in1k": {"input_size": 320},
+    "mambaout_tiny.in1k": {"input_size": 224},
+    "mambaout_base.in1k": {"input_size": 224},
+    "focalnet_tiny_lrf.ms_in1k": {"input_size": 224},
+    "focalnet_base_lrf.ms_in1k": {"input_size": 224},
+    "edgenext_xx_small.in1k": {"input_size": 256},
+    "edgenext_base.in21k_ft_in1k": {"input_size": 256},
+    "swin_tiny_patch4_window7_224": {"input_size": 224},
+    "swin_base_patch4_window7_224": {"input_size": 224},
+}
+
+
+def get_default_input_size(architecture_name):
+    """Get default input size for an architecture."""
+    return ARCHITECTURES.get(architecture_name, {}).get("input_size", 224)

model_cgd.py

@@ -0,0 +1,225 @@
+"""
+Circular Gaussian Distribution (CGD) for Image Orientation Estimation (Inference Only)
+
+Represents angles as probability distributions over discretized angle bins.
+Model output: Probability distribution over 360 angle bins (1 degree resolution)
+"""
+
+import math
+from typing import Dict, Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms as transforms
+import pytorch_lightning as pl
+import timm
+import timm.data
+from PIL import Image
+import numpy as np
+from loguru import logger
+
+
+class CircularGaussianDistribution(nn.Module):
+    """Circular Gaussian Distribution module for 360 degree image orientation."""
+
+    def __init__(self, num_bins: int = 360, sigma: float = 6.0):
+        super().__init__()
+        self.num_bins = num_bins
+        self.sigma = sigma
+        self.bin_size = 360.0 / num_bins
+
+        bin_centers = torch.arange(0, 360, self.bin_size)
+        self.register_buffer('bin_centers', bin_centers)
+
+        logger.info(f"CGD: {num_bins} bins, range [0, 360), sigma={sigma}")
+
+    def distribution_to_angle(self, distributions: torch.Tensor, method: str = 'argmax') -> torch.Tensor:
+        """Extract angles from probability distributions.
+
+        Args:
+            distributions: Probability distributions [B, num_bins]
+            method: 'argmax', 'weighted_average', or 'peak_fitting'
+
+        Returns:
+            angles: Extracted angles in degrees [B] in [0, 360)
+        """
+        if method == 'argmax':
+            peak_indices = torch.argmax(distributions, dim=1)
+            angles = self.bin_centers[peak_indices]
+
+        elif method == 'weighted_average':
+            weights = distributions / (distributions.sum(dim=1, keepdim=True) + 1e-8)
+            bin_angles_rad = self.bin_centers * torch.pi / 180.0
+            cos_components = torch.cos(bin_angles_rad)
+            sin_components = torch.sin(bin_angles_rad)
+            avg_cos = torch.sum(weights * cos_components.unsqueeze(0), dim=1)
+            avg_sin = torch.sum(weights * sin_components.unsqueeze(0), dim=1)
+            angles = torch.atan2(avg_sin, avg_cos) * 180.0 / torch.pi
+            angles = angles % 360.0
+
+        elif method == 'peak_fitting':
+            peak_indices = torch.argmax(distributions, dim=1)
+            angles = torch.zeros_like(peak_indices, dtype=torch.float)
+            for i in range(distributions.shape[0]):
+                peak_idx = peak_indices[i].item()
+                if 0 < peak_idx < self.num_bins - 1:
+                    y1 = distributions[i, peak_idx - 1]
+                    y2 = distributions[i, peak_idx]
+                    y3 = distributions[i, peak_idx + 1]
+                    a = 0.5 * (y1 - 2*y2 + y3)
+                    b = 0.5 * (y3 - y1)
+                    if abs(a) > 1e-8:
+                        offset = -b / (2 * a)
+                        offset = torch.clamp(offset, -0.5, 0.5)
+                    else:
+                        offset = 0
+                    angles[i] = self.bin_centers[peak_idx] + offset * self.bin_size
+                else:
+                    angles[i] = self.bin_centers[peak_idx]
+        else:
+            raise ValueError(f"Unknown extraction method: {method}")
+
+        angles = angles % 360.0
+        return angles
+
+    def get_distribution_uncertainty(self, distributions: torch.Tensor) -> torch.Tensor:
+        """Calculate entropy-based uncertainty from distribution."""
+        log_probs = torch.log(distributions + 1e-8)
+        entropy = -torch.sum(distributions * log_probs, dim=1)
+        max_entropy = math.log(self.num_bins)
+        return entropy / max_entropy
+
+
+class CGDAngleEstimation(pl.LightningModule):
+    """CGD model for 360 degree image orientation estimation (inference only)."""
+
+    def __init__(
+        self,
+        batch_size: int = 16,
+        train_dir: str = "",
+        model_name: str = "vit_tiny_patch16_224",
+        learning_rate: float = 0.001,
+        validation_split: float = 0.1,
+        random_seed: int = 42,
+        image_size: int = 224,
+        num_bins: int = 360,
+        sigma: float = 6.0,
+        inference_method: str = 'argmax',
+        loss_type: str = 'kl_divergence',
+        test_dir=None,
+        test_rotation_range=360.0,
+        test_random_seed=42,
+    ) -> None:
+        super().__init__()
+        self.save_hyperparameters()
+
+        self.model_name = model_name
+        self.learning_rate = learning_rate
+        self.batch_size = batch_size
+        self.train_dir = train_dir
+        self.validation_split = validation_split
+        self.random_seed = random_seed
+        self.image_size = image_size
+        self.num_bins = num_bins
+        self.sigma = sigma
+        self.inference_method = inference_method
+        self.loss_type = loss_type
+
+        self.model = timm.create_model(model_name, pretrained=True, num_classes=num_bins)
+        self.cgd = CircularGaussianDistribution(num_bins=num_bins, sigma=sigma)
+
+    @classmethod
+    def try_load(cls, checkpoint_path=None, **kwargs):
+        """Load model from checkpoint."""
+        if checkpoint_path:
+            logger.info(f"Loading model from checkpoint: {checkpoint_path}")
+            model = cls.load_from_checkpoint(checkpoint_path, **kwargs)
+            logger.info("Model loaded successfully from checkpoint")
+            return model
+        raise FileNotFoundError("Checkpoint file not found")
+
+    @classmethod
+    def from_pretrained(cls, repo_id, model_name=None):
+        """Load a pretrained model from HuggingFace Hub.
+
+        Args:
+            repo_id: HuggingFace repo ID (e.g. "maxwoe/image-rotation-angle-estimation")
+            model_name: Display name or checkpoint filename from config.json.
+                Defaults to the default model.
+        """
+        import json
+        from huggingface_hub import hf_hub_download
+
+        config_path = hf_hub_download(repo_id=repo_id, filename="config.json")
+        with open(config_path) as f:
+            config = json.load(f)
+
+        if model_name is None:
+            model_name = config["default_model"]
+
+        # Look up by display name or by filename
+        if model_name in config["models"]:
+            model_info = config["models"][model_name]
+        else:
+            model_info = None
+            for info in config["models"].values():
+                if info["filename"] == model_name:
+                    model_info = info
+                    break
+            if model_info is None:
+                available = [i["filename"] for i in config["models"].values()]
+                raise ValueError(f"Unknown model: {model_name}. Available: {available}")
+
+        ckpt_path = hf_hub_download(repo_id=repo_id, filename=model_info["filename"])
+        model = cls.try_load(checkpoint_path=ckpt_path, image_size=model_info["input_size"])
+        model.eval()
+        return model
+
+    def forward(self, x: torch.Tensor, return_logits: bool = False) -> torch.Tensor:
+        """Forward pass returning probability distribution over angles."""
+        logits = self.model(x)
+        if return_logits:
+            return logits
+        return F.softmax(logits, dim=1)
+
+    def predict_angle(self, image) -> float:
+        """Detect the current orientation angle of an image.
+
+        Args:
+            image: PIL Image, numpy array, or file path string.
+                   For best results, pass PIL Image or numpy array directly.
+
+        Returns:
+            Predicted rotation angle in degrees [0, 360).
+        """
+        self.eval()
+
+        if isinstance(image, str):
+            image = Image.open(image).convert('RGB')
+        elif isinstance(image, np.ndarray):
+            image = Image.fromarray(image).convert('RGB')
+        elif not isinstance(image, Image.Image):
+            raise TypeError(f"Expected PIL Image, numpy array, or file path, got {type(image)}")
+        else:
+            image = image.convert('RGB')
+
+        try:
+            data_config = timm.data.resolve_model_data_config(self.hparams.model_name)
+            data_config['crop_pct'] = 1.0
+            data_config['input_size'] = (3, self.image_size, self.image_size)
+            transform = timm.data.create_transform(**data_config, is_training=False)
+        except Exception:
+            transform = transforms.Compose([
+                transforms.Resize((self.image_size, self.image_size)),
+                transforms.ToTensor(),
+                transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+            ])
+
+        image_tensor = transform(image).unsqueeze(0)
+
+        with torch.no_grad():
+            pred_distributions = self(image_tensor)
+            angle = self.cgd.distribution_to_angle(pred_distributions, method=self.inference_method).item()
+
+        return angle

requirements.txt

@@ -0,0 +1,10 @@
+torch>=2.1.0
+torchvision>=0.16.0
+pytorch-lightning>=2.1.0
+timm>=0.9.0
+gradio>=4.0.0
+Pillow>=10.0.0
+loguru>=0.7.0
+numpy>=1.24.0
+opencv-python>=4.5.0
+huggingface_hub>=0.20.0

rotation_utils.py

@@ -0,0 +1,71 @@
+"""Rotation utilities that preserve image content without artificial borders."""
+
+import cv2
+import numpy as np
+import math
+
+
+def rotate_image(img, angle, rotation_center=None, expand=False, border_mode=cv2.BORDER_CONSTANT, border_value=0):
+    """Rotates an image (angle in degrees) and optionally expands to avoid cropping."""
+    h, w = img.shape[:2]
+    if rotation_center is None:
+        rotation_center = (w/2, h/2)
+
+    M = cv2.getRotationMatrix2D(rotation_center, angle, 1.0)
+
+    if expand:
+        abs_cos = abs(M[0, 0])
+        abs_sin = abs(M[0, 1])
+        wn = int(h * abs_sin + w * abs_cos)
+        hn = int(h * abs_cos + w * abs_sin)
+        M[0, 2] += wn/2 - rotation_center[0]
+        M[1, 2] += hn/2 - rotation_center[1]
+    else:
+        wn, hn = w, h
+
+    rotated = cv2.warpAffine(
+        img, M, (wn, hn), borderMode=border_mode, borderValue=border_value)
+
+    return rotated, M
+
+
+def largest_rotated_rect(w, h, angle):
+    """Compute the largest axis-aligned rectangle within a rotated rectangle."""
+    if w <= 0 or h <= 0:
+        return 0, 0
+
+    width_is_longer = w >= h
+    side_long, side_short = (w, h) if width_is_longer else (h, w)
+
+    sin_a, cos_a = abs(math.sin(angle)), abs(math.cos(angle))
+    if side_short <= 2.*sin_a*cos_a*side_long or abs(sin_a-cos_a) < 1e-10:
+        x = 0.5*side_short
+        wr, hr = (x/sin_a, x/cos_a) if width_is_longer else (x/cos_a, x/sin_a)
+    else:
+        cos_2a = cos_a*cos_a - sin_a*sin_a
+        wr, hr = (w*cos_a - h*sin_a)/cos_2a, (h*cos_a - w*sin_a)/cos_2a
+
+    return wr, hr
+
+
+def rotate_image_crop_max_area(image, angle):
+    """Rotate image and crop to the largest inscribed rectangle (no borders).
+
+    Args:
+        image: numpy array (OpenCV image)
+        angle: Rotation angle in degrees
+
+    Returns:
+        Rotated and cropped numpy array
+    """
+    h, w = image.shape[:2]
+    rotated, _ = rotate_image(image, angle, expand=True)
+    wr, hr = largest_rotated_rect(w, h, math.radians(angle))
+
+    h_rot, w_rot = rotated.shape[:2]
+    y1 = h_rot//2 - int(hr/2)
+    y2 = y1 + int(hr)
+    x1 = w_rot//2 - int(wr/2)
+    x2 = x1 + int(wr)
+
+    return rotated[y1:y2, x1:x2]