""" 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()