| 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): |
| state = torch.load(path, map_location=self.device, weights_only=True) |
| if "model_state_dict" in state: |
| self.model.load_state_dict(state["model_state_dict"]) |
| else: |
| self.model.load_state_dict(state) |
| 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)") |
|
|
| @torch.no_grad() |
| 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] |
|
|
| 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": { |
| "low_frequency": freq_magnitudes[0], |
| "mid_frequency": freq_magnitudes[1], |
| "high_frequency": freq_magnitudes[2], |
| }, |
| "processing_time_ms": elapsed, |
| } |
| return result |
|
|
| def _preprocess(self, image: Image.Image) -> torch.Tensor: |
| image = image.convert("RGB").resize( |
| (self.image_size, self.image_size), Image.Resampling.LANCZOS) |
| 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) |
|
|
