File size: 3,601 Bytes
abeae79 | 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 | 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)
|