maxwoe/image-rotation-angle-estimation

1"""HuggingFace Spaces demo for Image Rotation Angle Estimation.
2
3Two-step interactive demo:
41. Upload an image and click "Random Rotate" to apply a random rotation
52. Click "Correct Orientation" to see the model predict and correct the angle
6"""
7
8import json
9
10import gradio as gr
11import torch
12from PIL import Image
13import os
14import random
15import numpy as np
16from loguru import logger
17from huggingface_hub import hf_hub_download
18
19from model_cgd import CGDAngleEstimation
20from architectures import get_default_input_size
21from rotation_utils import rotate_image_crop_max_area
22
23# HuggingFace Hub configuration
24HF_REPO_ID = os.environ.get("HF_MODEL_REPO", "maxwoe/image-rotation-angle-estimation")
25
26# Fetch config.json from model repo (this is the HF download-tracking query file)
27_config_path = hf_hub_download(repo_id=HF_REPO_ID, filename="config.json")
28with open(_config_path) as _f:
29    _config = json.load(_f)
30HF_MODELS = _config["models"]
31HF_DEFAULT_MODEL = _config["default_model"]
32
33# Global model state
34model = None
35current_model_name = None
36
37
38def get_device():
39    return "cuda:0" if torch.cuda.is_available() else "cpu"
40
41
42def load_model(name):
43    """Download and load a model from HuggingFace Hub."""
44    global model, current_model_name
45    if name == current_model_name and model is not None:
46        return gr.Info(f"Model already loaded: {name}")
47
48    if name not in HF_MODELS:
49        return gr.Warning(f"Unknown model: {name}")
50
51    info = HF_MODELS[name]
52    logger.info(f"Downloading {info['filename']} from {HF_REPO_ID}...")
53    local_path = hf_hub_download(repo_id=HF_REPO_ID, filename=info["filename"])
54
55    architecture = info["architecture"]
56    image_size = get_default_input_size(architecture)
57
58    logger.info(f"Loading model from {local_path}...")
59    new_model = CGDAngleEstimation.try_load(checkpoint_path=local_path, image_size=image_size)
60    new_model.eval()
61
62    device = get_device()
63    if device.startswith("cuda"):
64        new_model = new_model.to(device)
65
66    model = new_model
67    current_model_name = name
68    logger.info(f"Model loaded: {name} on {device}")
69    return gr.Info(f"Loaded: {name} ({device})")
70
71
72def store_original(image):
73    """Store the uploaded image as the original for rotation."""
74    if image is None:
75        return None, ""
76    if isinstance(image, np.ndarray):
77        image = Image.fromarray(image)
78    return image, ""
79
80
81def random_rotate(original):
82    """Apply a random rotation to the original uploaded image."""
83    if original is None:
84        return None, None, ""
85
86    angle = random.uniform(0, 360)
87    img_array = np.array(original)
88    rotated_array = rotate_image_crop_max_area(img_array, angle)
89    rotated = Image.fromarray(rotated_array)
90    return rotated, angle, f"{angle:.1f}°"
91
92
93def correct_orientation(image):
94    """Predict the rotation angle and correct the image."""
95    if image is None:
96        return None, "Please upload and rotate an image first."
97    if model is None:
98        return None, "Model is still loading, please wait..."
99
100    if isinstance(image, np.ndarray):
101        image = Image.fromarray(image)
102
103    predicted_angle = model.predict_angle(image)
104
105    corrected = image.rotate(-predicted_angle, expand=True, fillcolor=(255, 255, 255))
106
107    return corrected, f"Predicted rotation: {predicted_angle:.2f}°"
108
109
110# Build UI
111app = gr.Blocks(title="Image Rotation Angle Estimation")
112with app:
113    gr.HTML("<h1>Image Rotation Angle Estimation</h1>")
114    gr.Markdown(
115        "Upload an image, apply a random rotation, and see the model predict and correct the angle.\n\n"
116        "Uses the **CGD** (Circular Gaussian Distribution) method with **MambaOut Base** architecture. "
117    )
118
119    original_image_state = gr.State(value=None)
120    actual_angle_state = gr.State(value=None)
121
122    model_dropdown = gr.Dropdown(
123        choices=list(HF_MODELS.keys()),
124        value=HF_DEFAULT_MODEL,
125        label="Model",
126    )
127    model_dropdown.change(load_model, inputs=[model_dropdown])
128
129    with gr.Row():
130        with gr.Column():
131            input_image = gr.Image(type="pil", label="Upload Image", height=400, format="png")
132            rotate_btn = gr.Button("Random Rotate", variant="secondary", size="lg")
133        with gr.Column():
134            corrected_image = gr.Image(label="Corrected Image", height=400, interactive=False, format="png")
135            correct_btn = gr.Button("Correct Orientation", variant="primary", size="lg")
136
137    with gr.Row():
138        rotation_info = gr.Textbox(label="Applied Rotation", lines=1, interactive=False)
139        result_text = gr.Textbox(label="Predicted Rotation", lines=1, interactive=False)
140
141    gr.Examples(
142        examples=[
143            "examples/COCO_val2014_000000168337.jpg",
144            "examples/COCO_val2014_000000122166.jpg",
145            "examples/COCO_val2014_000000446053.jpg",
146            "examples/COCO_val2014_000000477919.jpg",
147            "examples/COCO_val2014_000000016995.jpg",
148        ],
149        inputs=input_image,
150        fn=store_original,
151        outputs=[original_image_state, rotation_info],
152        cache_examples=False,
153        run_on_click=True,
154    )
155
156    input_image.upload(
157        store_original,
158        inputs=[input_image],
159        outputs=[original_image_state, rotation_info],
160    )
161    rotate_btn.click(
162        random_rotate,
163        inputs=[original_image_state],
164        outputs=[input_image, actual_angle_state, rotation_info],
165    )
166    correct_btn.click(
167        correct_orientation,
168        inputs=[input_image],
169        outputs=[corrected_image, result_text],
170    )
171
172    app.load(lambda: load_model(HF_DEFAULT_MODEL))
173
174app.launch()
175