Spaces:
Sleeping
Sleeping
File size: 5,646 Bytes
596aaa1 1d745c0 596aaa1 1d745c0 596aaa1 1d745c0 596aaa1 1d745c0 596aaa1 1d745c0 596aaa1 1d745c0 596aaa1 1d745c0 596aaa1 1d745c0 596aaa1 1d745c0 596aaa1 1d745c0 596aaa1 1d745c0 | 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 | import logging
import os
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn
from PIL import Image
from torchvision import transforms
from app.architecture import AdvancedBreastCancerModel
logger = logging.getLogger(__name__)
# ImageNet normalisation (same as SensiNet training pipeline)
TRANSFORM = transforms.Compose([
transforms.Resize((299, 299)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
WEIGHTS_DIR = Path(__file__).resolve().parent.parent / "weights"
DEFAULT_WEIGHTS = WEIGHTS_DIR / "advanced_model_best.pth"
# Malignancy probability threshold (same as SensiNet default)
THRESHOLD = 0.40
# Number of Bayesian MC-Dropout forward passes
MC_PASSES = 10
def _prob_to_birads(prob: float) -> int:
"""Map malignancy probability to BI-RADS category."""
if prob < 0.10:
return 1 # Negative
if prob < 0.25:
return 2 # Benign
if prob < 0.50:
return 3 # Probably benign
if prob < 0.75:
return 4 # Suspicious
return 5 # Highly suggestive of malignancy
def _birads_findings(birads: int, prob: float, prediction: str) -> str:
templates = {
1: "No suspicious findings detected. Mammographic appearance is unremarkable.",
2: "Benign-appearing pattern identified. Correlate with prior imaging if available.",
3: "Probably benign appearance. Short-interval follow-up may be considered.",
4: "Suspicious abnormality pattern detected. Tissue biopsy is recommended.",
5: "Highly suggestive of malignancy. Urgent diagnostic workup is recommended.",
}
base = templates.get(birads, "Analysis complete.")
return f"Model prediction: {prediction} (probability {prob:.1%}). {base}"
class MammogramModel:
"""Loads the SensiNet dual-stream model and runs inference."""
def __init__(self) -> None:
self.mode = os.getenv("MODEL_MODE", "real")
self.version = os.getenv("MODEL_VERSION", "sensinet-v1")
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self._model: AdvancedBreastCancerModel | None = None
weights_path = Path(os.getenv("MODEL_WEIGHTS", str(DEFAULT_WEIGHTS)))
if weights_path.exists():
self._load_model(weights_path)
else:
logger.warning("Weights not found at %s — falling back to mock mode", weights_path)
self.mode = "mock"
def _load_model(self, weights_path: Path) -> None:
logger.info("Loading SensiNet model from %s onto %s …", weights_path, self.device)
net = AdvancedBreastCancerModel()
state = torch.load(weights_path, map_location=self.device, weights_only=False)
net.load_state_dict(state)
net.to(self.device)
net.eval()
self._model = net
logger.info("Model loaded successfully.")
# ------------------------------------------------------------------
def predict(self, image: Image.Image) -> dict:
if self._model is None or self.mode == "mock":
return self._mock_predict(image)
return self._real_predict(image)
# ------------------------------------------------------------------
# Real inference with Bayesian MC-Dropout
# ------------------------------------------------------------------
def _real_predict(self, image: Image.Image) -> dict:
rgb = image.convert("RGB")
tensor = TRANSFORM(rgb).unsqueeze(0).to(self.device)
def enable_dropout(m: nn.Module) -> None:
if isinstance(m, (nn.Dropout, nn.Dropout2d)):
m.train()
self._model.apply(enable_dropout)
mc_predictions: list[float] = []
with torch.no_grad():
for _ in range(MC_PASSES):
logits = self._model(tensor)
prob = torch.sigmoid(logits).item()
mc_predictions.append(prob)
self._model.eval()
prob_malig = float(np.mean(mc_predictions))
variance = float(np.var(mc_predictions))
decision_confidence = max(0.50, 0.99 - (variance * 2.0))
if prob_malig < 0.10 or prob_malig > 0.90:
decision_confidence = min(0.99, decision_confidence + 0.10)
prediction = "Malignant" if prob_malig >= THRESHOLD else "Benign"
birads = _prob_to_birads(prob_malig)
return {
"birads": birads,
"confidence": round(decision_confidence, 3),
"malignancy_probability": round(prob_malig, 3),
"findings_text": _birads_findings(birads, prob_malig, prediction),
"model_version": self.version,
}
# ------------------------------------------------------------------
# Deterministic mock fallback (no weights needed)
# ------------------------------------------------------------------
@staticmethod
def _mock_predict(image: Image.Image) -> dict:
import hashlib
arr = np.array(image.convert("L"), dtype=np.float32) / 255.0
digest = hashlib.sha256(arr.tobytes()).hexdigest()
seed = int(digest[:8], 16)
rng = np.random.default_rng(seed)
raw = float(min(max(arr.mean() + rng.uniform(-0.04, 0.04), 0.0), 1.0))
birads = _prob_to_birads(raw)
return {
"birads": birads,
"confidence": round(max(0.55, min(0.98, 0.55 + abs(raw - 0.5))), 3),
"malignancy_probability": round(raw, 3),
"findings_text": _birads_findings(birads, raw, "Malignant" if raw >= THRESHOLD else "Benign"),
"model_version": "mock-v1",
}
|