Spaces:
Sleeping
Sleeping
File size: 7,316 Bytes
6cc8ae1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 | """
Inference service.
Loads the model on startup and provides prediction functionality.
"""
import sys
import time
from pathlib import Path
from typing import Optional
import torch
import torch.nn.functional as F
from PIL import Image
from backend.app.core.config import get_settings
from backend.app.core.logging import logger
class InferenceService:
"""Singleton service for model inference."""
def __init__(self):
self.model = None
self.class_names: list[str] = []
self.device = torch.device('cpu')
self.settings = get_settings()
self._loaded = False
def load(self) -> None:
"""Load model and class names on startup."""
if self._loaded:
return
logger.info("Loading model...")
start = time.time()
# Detect device
if torch.cuda.is_available():
self.device = torch.device('cuda')
elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
self.device = torch.device('mps')
# Load class names
classes_path = Path(self.settings.classes_path)
if classes_path.exists():
with open(classes_path, 'r') as f:
self.class_names = [c.strip() for c in f.read().split(',') if c.strip()]
# Load model
model_path = Path(self.settings.model_path)
if not model_path.exists():
logger.warning(f"Model file not found: {model_path}")
logger.info("Service will start without a model. Upload a model to enable predictions.")
return
try:
checkpoint = torch.load(model_path, map_location=self.device, weights_only=False)
if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
# State dict checkpoint from new training pipeline
model_name = checkpoint.get('model_name', self.settings.model_name)
if 'class_names' in checkpoint:
self.class_names = checkpoint['class_names']
self.model = self._create_model(model_name, len(self.class_names))
self.model.load_state_dict(checkpoint['model_state_dict'])
elif isinstance(checkpoint, torch.nn.Module):
# Full model (legacy format)
self.model = checkpoint
else:
# Try loading as full model directly
self.model = checkpoint
if isinstance(self.model, torch.nn.Module):
self.model = self.model.to(self.device)
self.model.eval()
self._loaded = True
elapsed = time.time() - start
logger.info(f"Model loaded in {elapsed:.2f}s on {self.device} ({len(self.class_names)} classes)")
except Exception as e:
logger.error(f"Failed to load model: {e}")
raise
def _create_model(self, model_name: str, num_classes: int) -> torch.nn.Module:
"""Create model architecture by name."""
# Add project root to path for ML imports
project_root = Path(__file__).resolve().parents[3]
if str(project_root) not in sys.path:
sys.path.insert(0, str(project_root))
if model_name in ('resnet', 'resnet50'):
from ml.src.models.resnet import CattleResNet
return CattleResNet(num_classes=num_classes, pretrained=False, freeze_backbone=False)
elif model_name == 'cnn':
from ml.src.models.cnn import CattleCNN
return CattleCNN(num_classes=num_classes)
elif model_name == 'mlp':
from ml.src.models.mlp import CattleMLP
return CattleMLP(num_classes=num_classes)
elif model_name in ('vit', 'vit_base'):
from ml.src.models.vit import CattleViT
return CattleViT(num_classes=num_classes, pretrained=False, freeze_backbone=False)
else:
raise ValueError(f"Unknown model type: {model_name}")
@torch.no_grad()
def predict(self, image: Image.Image, top_k: int = 3) -> dict:
"""
Run prediction on a PIL Image.
Returns prediction dict.
"""
if self.model is None:
# Return a mock prediction if model is missing to prevent UI crash
import random
# Use real class names if available, otherwise fallback
classes = self.class_names if self.class_names else [
"Gir", "Sahiwal", "Kankrej", "Tharparkar", "Red Sindhi"
]
mock_breed = random.choice(classes)
other_breeds = [c for c in classes if c != mock_breed]
mock_top_k = [{'breed': mock_breed, 'confidence': 0.875}]
if other_breeds:
mock_top_k.append({'breed': other_breeds[0], 'confidence': 0.082})
if len(other_breeds) > 1:
mock_top_k.append({'breed': other_breeds[1], 'confidence': 0.043})
return {
'predicted_breed': mock_breed,
'confidence': 0.875,
'top_k': mock_top_k,
'model_version': 'mock-v1.0 (No Model Loaded)',
'inference_time_ms': 12.5,
'warning': '⚠️ PyTorch model file missing from server. Displaying mock AI prediction.'
}
from torchvision import transforms
start = time.time()
# Preprocess
transform = transforms.Compose([
transforms.Resize((self.settings.image_size, self.settings.image_size)),
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225],
),
])
tensor = transform(image).unsqueeze(0).to(self.device)
# Inference
outputs = self.model(tensor)
probabilities = F.softmax(outputs, dim=1).squeeze()
# Top-k
top_k_val = min(top_k, len(self.class_names))
top_probs, top_indices = torch.topk(probabilities, top_k_val)
top_k_results = []
for prob, idx in zip(top_probs, top_indices):
breed_name = self.class_names[idx.item()] if idx.item() < len(self.class_names) else f"class_{idx.item()}"
top_k_results.append({
'breed': breed_name,
'confidence': round(prob.item(), 4),
})
best = top_k_results[0]
elapsed_ms = (time.time() - start) * 1000
result = {
'predicted_breed': best['breed'],
'confidence': best['confidence'],
'top_k': top_k_results,
'model_version': self.settings.model_version,
'inference_time_ms': round(elapsed_ms, 2),
}
# Confidence warning
if best['confidence'] < 0.5:
result['warning'] = 'Prediction confidence is low. Use a clear side-view image for better results.'
elif best['confidence'] < 0.75:
result['warning'] = 'Prediction confidence is moderate. Use a clear side-view image for better results.'
return result
@property
def is_loaded(self) -> bool:
return self._loaded and self.model is not None
# Singleton
inference_service = InferenceService()
|