maxwoe/image-rotation-angle-estimation
0
1"""
2Circular Gaussian Distribution (CGD) for Image Orientation Estimation (Inference Only)
3
4Represents angles as probability distributions over discretized angle bins.
5Model output: Probability distribution over 360 angle bins (1 degree resolution)
6"""
7
8import math
9from typing import Dict, Any
10
11import torch
12import torch.nn as nn
13import torch.nn.functional as F
14import torchvision.transforms as transforms
15import pytorch_lightning as pl
16import timm
17import timm.data
18from PIL import Image
19import numpy as np
20from loguru import logger
21
22
23class CircularGaussianDistribution(nn.Module):
24 """Circular Gaussian Distribution module for 360 degree image orientation."""
25
26 def __init__(self, num_bins: int = 360, sigma: float = 6.0):
27 super().__init__()
28 self.num_bins = num_bins
29 self.sigma = sigma
30 self.bin_size = 360.0 / num_bins
31
32 bin_centers = torch.arange(0, 360, self.bin_size)
33 self.register_buffer('bin_centers', bin_centers)
34
35 logger.info(f"CGD: {num_bins} bins, range [0, 360), sigma={sigma}")
36
37 def distribution_to_angle(self, distributions: torch.Tensor, method: str = 'argmax') -> torch.Tensor:
38 """Extract angles from probability distributions.
39
40 Args:
41 distributions: Probability distributions [B, num_bins]
42 method: 'argmax', 'weighted_average', or 'peak_fitting'
43
44 Returns:
45 angles: Extracted angles in degrees [B] in [0, 360)
46 """
47 if method == 'argmax':
48 peak_indices = torch.argmax(distributions, dim=1)
49 angles = self.bin_centers[peak_indices]
50
51 elif method == 'weighted_average':
52 weights = distributions / (distributions.sum(dim=1, keepdim=True) + 1e-8)
53 bin_angles_rad = self.bin_centers * torch.pi / 180.0
54 cos_components = torch.cos(bin_angles_rad)
55 sin_components = torch.sin(bin_angles_rad)
56 avg_cos = torch.sum(weights * cos_components.unsqueeze(0), dim=1)
57 avg_sin = torch.sum(weights * sin_components.unsqueeze(0), dim=1)
58 angles = torch.atan2(avg_sin, avg_cos) * 180.0 / torch.pi
59 angles = angles % 360.0
60
61 elif method == 'peak_fitting':
62 peak_indices = torch.argmax(distributions, dim=1)
63 angles = torch.zeros_like(peak_indices, dtype=torch.float)
64 for i in range(distributions.shape[0]):
65 peak_idx = peak_indices[i].item()
66 if 0 < peak_idx < self.num_bins - 1:
67 y1 = distributions[i, peak_idx - 1]
68 y2 = distributions[i, peak_idx]
69 y3 = distributions[i, peak_idx + 1]
70 a = 0.5 * (y1 - 2*y2 + y3)
71 b = 0.5 * (y3 - y1)
72 if abs(a) > 1e-8:
73 offset = -b / (2 * a)
74 offset = torch.clamp(offset, -0.5, 0.5)
75 else:
76 offset = 0
77 angles[i] = self.bin_centers[peak_idx] + offset * self.bin_size
78 else:
79 angles[i] = self.bin_centers[peak_idx]
80 else:
81 raise ValueError(f"Unknown extraction method: {method}")
82
83 angles = angles % 360.0
84 return angles
85
86 def get_distribution_uncertainty(self, distributions: torch.Tensor) -> torch.Tensor:
87 """Calculate entropy-based uncertainty from distribution."""
88 log_probs = torch.log(distributions + 1e-8)
89 entropy = -torch.sum(distributions * log_probs, dim=1)
90 max_entropy = math.log(self.num_bins)
91 return entropy / max_entropy
92
93
94class CGDAngleEstimation(pl.LightningModule):
95 """CGD model for 360 degree image orientation estimation (inference only)."""
96
97 def __init__(
98 self,
99 batch_size: int = 16,
100 train_dir: str = "",
101 model_name: str = "vit_tiny_patch16_224",
102 learning_rate: float = 0.001,
103 validation_split: float = 0.1,
104 random_seed: int = 42,
105 image_size: int = 224,
106 num_bins: int = 360,
107 sigma: float = 6.0,
108 inference_method: str = 'argmax',
109 loss_type: str = 'kl_divergence',
110 test_dir=None,
111 test_rotation_range=360.0,
112 test_random_seed=42,
113 ) -> None:
114 super().__init__()
115 self.save_hyperparameters()
116
117 self.model_name = model_name
118 self.learning_rate = learning_rate
119 self.batch_size = batch_size
120 self.train_dir = train_dir
121 self.validation_split = validation_split
122 self.random_seed = random_seed
123 self.image_size = image_size
124 self.num_bins = num_bins
125 self.sigma = sigma
126 self.inference_method = inference_method
127 self.loss_type = loss_type
128
129 self.model = timm.create_model(model_name, pretrained=True, num_classes=num_bins)
130 self.cgd = CircularGaussianDistribution(num_bins=num_bins, sigma=sigma)
131
132 @classmethod
133 def try_load(cls, checkpoint_path=None, **kwargs):
134 """Load model from checkpoint."""
135 if checkpoint_path:
136 logger.info(f"Loading model from checkpoint: {checkpoint_path}")
137 model = cls.load_from_checkpoint(checkpoint_path, **kwargs)
138 logger.info("Model loaded successfully from checkpoint")
139 return model
140 raise FileNotFoundError("Checkpoint file not found")
141
142 @classmethod
143 def from_pretrained(cls, repo_id, model_name=None):
144 """Load a pretrained model from HuggingFace Hub.
145
146 Args:
147 repo_id: HuggingFace repo ID (e.g. "maxwoe/image-rotation-angle-estimation")
148 model_name: Display name or checkpoint filename from config.json.
149 Defaults to the default model.
150 """
151 import json
152 from huggingface_hub import hf_hub_download
153
154 config_path = hf_hub_download(repo_id=repo_id, filename="config.json")
155 with open(config_path) as f:
156 config = json.load(f)
157
158 if model_name is None:
159 model_name = config["default_model"]
160
161 # Look up by display name or by filename
162 if model_name in config["models"]:
163 model_info = config["models"][model_name]
164 else:
165 model_info = None
166 for info in config["models"].values():
167 if info["filename"] == model_name:
168 model_info = info
169 break
170 if model_info is None:
171 available = [i["filename"] for i in config["models"].values()]
172 raise ValueError(f"Unknown model: {model_name}. Available: {available}")
173
174 ckpt_path = hf_hub_download(repo_id=repo_id, filename=model_info["filename"])
175 model = cls.try_load(checkpoint_path=ckpt_path, image_size=model_info["input_size"])
176 model.eval()
177 return model
178
179 def forward(self, x: torch.Tensor, return_logits: bool = False) -> torch.Tensor:
180 """Forward pass returning probability distribution over angles."""
181 logits = self.model(x)
182 if return_logits:
183 return logits
184 return F.softmax(logits, dim=1)
185
186 def predict_angle(self, image) -> float:
187 """Detect the current orientation angle of an image.
188
189 Args:
190 image: PIL Image, numpy array, or file path string.
191 For best results, pass PIL Image or numpy array directly.
192
193 Returns:
194 Predicted rotation angle in degrees [0, 360).
195 """
196 self.eval()
197
198 if isinstance(image, str):
199 image = Image.open(image).convert('RGB')
200 elif isinstance(image, np.ndarray):
201 image = Image.fromarray(image).convert('RGB')
202 elif not isinstance(image, Image.Image):
203 raise TypeError(f"Expected PIL Image, numpy array, or file path, got {type(image)}")
204 else:
205 image = image.convert('RGB')
206
207 try:
208 data_config = timm.data.resolve_model_data_config(self.hparams.model_name)
209 data_config['crop_pct'] = 1.0
210 data_config['input_size'] = (3, self.image_size, self.image_size)
211 transform = timm.data.create_transform(**data_config, is_training=False)
212 except Exception:
213 transform = transforms.Compose([
214 transforms.Resize((self.image_size, self.image_size)),
215 transforms.ToTensor(),
216 transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
217 ])
218
219 image_tensor = transform(image).unsqueeze(0)
220
221 with torch.no_grad():
222 pred_distributions = self(image_tensor)
223 angle = self.cgd.distribution_to_angle(pred_distributions, method=self.inference_method).item()
224
225 return angle
226