import base64 import io from typing import Tuple import numpy as np from PIL import Image import torch import torch.nn.functional as F from model import DEVICE, get_unet # ---------- Image loading / preprocessing ---------- def load_image_from_bytes(data: bytes, target_size: Tuple[int, int] = (240, 240)) -> np.ndarray: """ Load uploaded image (PNG/JPEG), resize to 240x240. Returns float32 (H, W) in original pixel range (0-255). """ img = Image.open(io.BytesIO(data)).convert("L") # grayscale img = img.resize(target_size, resample=Image.BILINEAR) arr = np.array(img).astype(np.float32) # (H, W), 0..255 return arr def build_4channel_input(arr: np.ndarray) -> np.ndarray: """ Model expects (H, W, 4) channels. We synthesize 4 'pseudo-modalities' from a single grayscale slice: ch0: z-score of original ch1: z-score of gamma 0.5 (brightens) ch2: z-score of gamma 2.0 (darkens) ch3: z-score of local min-max (CLAHE-like).[file:1] """ def zscore(x): m, s = x.mean(), x.std() + 1e-8 return (x - m) / s x = arr / 255.0 # [0, 1] ch0 = zscore(x) ch1 = zscore(np.power(np.clip(x, 0, 1), 0.5)) ch2 = zscore(np.power(np.clip(x, 0, 1), 2.0)) # local min-max in 8x8 patches x_local = x.copy() h, w = x_local.shape bs = 8 for i in range(0, h, bs): for j in range(0, w, bs): blk = x_local[i:i + bs, j:j + bs] mn, mx = blk.min(), blk.max() x_local[i:i + bs, j:j + bs] = (blk - mn) / (mx - mn + 1e-8) ch3 = zscore(x_local) img4 = np.stack([ch0, ch1, ch2, ch3], axis=-1) # (H, W, 4) return img4.astype(np.float32) # ---------- Heuristic survival from mask (new model) ---------- SURV_LABELS = { 0: "Short Survival", 1: "Mid Survival", 2: "Long Survival", } def heuristic_survival_from_masks(pred_masks: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: """ pred_masks: (B, H, W) tensor with classes {0 BG, 1 NCR/NET, 2 Edema, 3 ET}.[file:1] Returns: surv_classes: (B,) long tensor in {0,1,2} et_fracs : (B,) float tensor of ET fraction per slice """ batch_size = pred_masks.shape[0] surv_classes = [] et_fracs = [] for i in range(batch_size): mask = pred_masks[i] total_pixels = mask.numel() et_pixels = (mask == 3).sum().float() edema_pixels = (mask == 2).sum().float() ncr_pixels = (mask == 1).sum().float() et_frac = et_pixels / total_pixels tumor_frac = (et_pixels + edema_pixels + ncr_pixels) / total_pixels # Heuristic from notebook: larger ET / tumour burden → poorer survival.[file:1] if (et_frac >= 0.03) or (tumor_frac >= 0.12): surv_class = 0 # Short Survival elif (et_frac >= 0.01) or (tumor_frac >= 0.05): surv_class = 1 # Mid Survival else: surv_class = 2 # Long Survival surv_classes.append(surv_class) et_fracs.append(et_frac) surv_classes = torch.tensor(surv_classes, dtype=torch.long, device=pred_masks.device) et_fracs = torch.stack(et_fracs).to(pred_masks.device) return surv_classes, et_fracs # ---------- Core inference: segmentation + survival ---------- def run_segmentation_and_survival(arr: np.ndarray, patient_age: float): """ arr: (H, W) float32 in 0..255 patient_age: kept for API compatibility but not used by the heuristic.[file:1] Returns ------- mask : (H, W) uint8, labels 0..3 conf_map : (H, W) float32, max softmax prob per pixel surv_probs : (3,) float32, one-hot probabilities [short, medium, long] surv_class : scalar int, argmax of surv_probs """ unet = get_unet() # ---- segmentation ---- img4 = build_4channel_input(arr) # (H, W, 4) x = np.transpose(img4, (2, 0, 1)) # (4, H, W) x = torch.from_numpy(x).unsqueeze(0).to(DEVICE) # (1, 4, H, W) with torch.no_grad(), torch.amp.autocast( "cuda", enabled=(DEVICE.type == "cuda") ): logits = unet(x) # (1, 4, H, W) probs = torch.softmax(logits, dim=1) preds = torch.argmax(probs, dim=1) # (1, H, W) conf = probs.max(dim=1)[0] # (1, H, W) max prob # numpy outputs for downstream mask = preds[0].cpu().numpy().astype(np.uint8) conf_map = conf[0].cpu().numpy().astype(np.float32) # ---- heuristic survival from predicted mask ---- surv_classes, et_fracs = heuristic_survival_from_masks(preds) surv_idx = int(surv_classes[0].item()) # we don't have calibrated probabilities; return one-hot surv_probs = np.zeros(3, dtype=np.float32) surv_probs[surv_idx] = 1.0 return mask, conf_map, surv_probs, surv_idx # ---------- Visualisation helpers ---------- def colorize_mask(mask: np.ndarray) -> Image.Image: """ mask: (H, W) labels 0..3 Colors: 1: Necrotic core (NCR/NET) -> yellow 2: Edema -> green 3: Effective Tumor (ET) -> red """ h, w = mask.shape overlay = np.zeros((h, w, 4), dtype=np.uint8) colors = { 1: (255, 255, 0, 180), # class 1: necrotic core - yellow 2: ( 0, 255, 0, 180), # class 2: edema - green 3: (255, 0, 0, 200), # class 3: Effective Tumor (ET) - red } for cls, (r, g, b, a) in colors.items(): overlay[mask == cls] = np.array([r, g, b, a], dtype=np.uint8) return Image.fromarray(overlay, mode="RGBA") def blend_overlay(arr: np.ndarray, overlay_rgba: Image.Image) -> Image.Image: """ arr: (H, W) float32, any range (auto-rescaled to 0..255) overlay_rgba: RGBA PIL Image Returns: RGB PIL Image — grayscale MRI with colored tumor overlay. """ mn, mx = arr.min(), arr.max() display = ((arr - mn) / (mx - mn + 1e-8) * 255.0).astype(np.uint8) base_img = Image.fromarray(display, mode="L").convert("RGBA") blended = Image.alpha_composite(base_img, overlay_rgba) return blended.convert("RGB") def pil_to_base64(img: Image.Image) -> str: buff = io.BytesIO() img.save(buff, format="PNG") return "data:image/png;base64," + base64.b64encode(buff.getvalue()).decode("utf-8") # ---------- Stats ---------- def get_tumor_stats(mask: np.ndarray, conf_map: np.ndarray) -> dict: total = mask.size tumor_pixels = int((mask > 0).sum()) tumor_pct = round(tumor_pixels / total * 100.0, 2) class_counts = { "necrotic_core": int((mask == 1).sum()), "edema": int((mask == 2).sum()), "effective_tumor": int((mask == 3).sum()), # was "et" / "enhancing_tumor" } if tumor_pixels > 0: mean_conf = float(conf_map[mask > 0].mean()) else: mean_conf = 0.0 return { "tumor_detected": tumor_pixels > 0, "tumor_pixel_count": tumor_pixels, "tumor_area_pct": tumor_pct, "class_counts": class_counts, "mean_tumor_confidence": mean_conf, }