Spaces:
Paused
Paused
| import time | |
| import torch | |
| import torch.nn.functional as F | |
| import numpy as np | |
| from PIL import Image | |
| from pathlib import Path | |
| from typing import Optional, Dict, Any | |
| import sys | |
| sys.path.insert(0, str(Path(__file__).parent.parent / "model" / "src")) | |
| from model import MFFTWithExplainability, build_mfft | |
| class ModelServer: | |
| """ | |
| Wraps the MFFT model for production inference. | |
| Handles model loading, preprocessing, prediction, and warmup. | |
| """ | |
| def __init__(self, checkpoint_path: Optional[str] = None, variant: str = "base", | |
| image_size: int = 384): | |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| self.variant = variant | |
| self.image_size = image_size | |
| print(f"[ModelServer] Device: {self.device} | Variant: {variant} | Size: {image_size}") | |
| self.model = build_mfft(variant) | |
| self.model = self.model.to(self.device) | |
| self.is_loaded = False | |
| if checkpoint_path and Path(checkpoint_path).exists(): | |
| self.load_checkpoint(checkpoint_path) | |
| self.is_loaded = True | |
| print(f"[ModelServer] Model loaded from {checkpoint_path}") | |
| else: | |
| print(f"[ModelServer] No checkpoint found at {checkpoint_path}") | |
| print(f"[ModelServer] Running with untrained weights (random initialization)") | |
| self.warmup() | |
| def load_checkpoint(self, path: str): | |
| try: | |
| state = torch.load(path, map_location=self.device, weights_only=True) | |
| except Exception: | |
| state = torch.load(path, map_location=self.device, weights_only=False) | |
| raw_state = state["model_state_dict"] if isinstance(state, dict) and "model_state_dict" in state else state | |
| # Remap keys between cluster training pipeline and local model.py naming conventions | |
| cleaned_state = {} | |
| for k, v in raw_state.items(): | |
| nk = k.replace(".stages.", ".blocks.").replace("fga.gate.", "freq_guided_attn.freq_gate.") | |
| if "num_batches_tracked" in nk: | |
| continue | |
| cleaned_state[nk] = v | |
| self.model.load_state_dict(cleaned_state, strict=False) | |
| self.is_loaded = True | |
| def warmup(self, n_warmup: int = 3): | |
| dummy = torch.randn(1, 3, self.image_size, self.image_size).to(self.device) | |
| self.model.eval() | |
| with torch.no_grad(): | |
| for _ in range(n_warmup): | |
| _ = self.model(dummy) | |
| print(f"[ModelServer] Warmup complete ({n_warmup} iterations)") | |
| def predict(self, image: Image.Image) -> Dict[str, Any]: | |
| start = time.time() | |
| input_tensor = self._preprocess(image) | |
| logits, heatmaps = self.model(input_tensor, return_heatmap=True) | |
| probs = F.softmax(logits, dim=-1) | |
| pred = torch.argmax(probs, dim=-1).item() | |
| bands = self.model.decomposer(input_tensor) | |
| freq_magnitudes = [band.abs().mean().item() for band in bands] | |
| if len(freq_magnitudes) == 4: | |
| band_contribs = { | |
| "low_frequency": freq_magnitudes[0], | |
| "low_mid_frequency": freq_magnitudes[1], | |
| "mid_high_frequency": freq_magnitudes[2], | |
| "high_frequency": freq_magnitudes[3], | |
| } | |
| else: | |
| band_contribs = { | |
| "low_frequency": freq_magnitudes[0], | |
| "mid_frequency": freq_magnitudes[1] if len(freq_magnitudes) > 1 else 0.0, | |
| "high_frequency": freq_magnitudes[2] if len(freq_magnitudes) > 2 else 0.0, | |
| } | |
| elapsed = (time.time() - start) * 1000 | |
| result = { | |
| "prediction": pred, | |
| "real_prob": probs[0, 0].item(), | |
| "ai_prob": probs[0, 1].item(), | |
| "confidence": probs.max(dim=-1).values.item(), | |
| "heatmaps": heatmaps, | |
| "frequency_band_contributions": band_contribs, | |
| "processing_time_ms": elapsed, | |
| } | |
| return result | |
| def _preprocess(self, image: Image.Image) -> torch.Tensor: | |
| # Aspect-ratio preserving resize + center crop (matching training/eval pipeline) | |
| w, h = image.size | |
| target = self.image_size + 16 | |
| if w < h: | |
| new_w, new_h = target, max(target, int(h * target / w)) | |
| else: | |
| new_h, new_w = target, max(target, int(w * target / h)) | |
| image = image.convert("RGB").resize((new_w, new_h), Image.Resampling.BILINEAR) | |
| left = max(0, (new_w - self.image_size) // 2) | |
| top = max(0, (new_h - self.image_size) // 2) | |
| image = image.crop((left, top, left + self.image_size, top + self.image_size)) | |
| img_array = np.array(image, dtype=np.float32) / 255.0 | |
| mean = np.array([0.485, 0.456, 0.406], dtype=np.float32) | |
| std = np.array([0.229, 0.224, 0.225], dtype=np.float32) | |
| img_array = ((img_array - mean) / std).astype(np.float32) | |
| tensor = torch.from_numpy(img_array).permute(2, 0, 1).unsqueeze(0) | |
| return tensor.to(self.device) | |