""" VAE inference: mask (PIL) -> preprocess -> encode -> decode -> return one slice. Input: Single grayscale mask (any size). Preprocess: Grayscale, Resize(256,256), ToTensor [0,1], duplicate to 4 slices -> (1, 4, 256, 256) batched format. Output: decode(z) shape (1, 4, 256, 256). We return one slice (default index 2) as PNG bytes. """ import io import logging import os from typing import Optional, Tuple import cv2 import numpy as np import torch import torch.nn as nn import torchvision.transforms as T from PIL import Image from huggingface_hub import hf_hub_download from model import VAE logging.basicConfig(level=os.environ.get("LOG_LEVEL", "INFO").upper()) logger = logging.getLogger(__name__) # --- Config (override via env on Hugging Face) --- DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") MODEL_REPO = os.environ.get("MODEL_REPO", "tan200224/Synthetic-CT-Scan_VAE_Conditional") MODEL_FILENAME = os.environ.get("MODEL_FILENAME", "mask2pic_64model_47.pt") INPUT_SIZE = 256 OUTPUT_SLICE_INDEX = 2 # Which of the 4 slices to return (0..3) _model: Optional[nn.Module] = None def load_model(repo_id: str = MODEL_REPO, filename: str = MODEL_FILENAME) -> nn.Module: """Download checkpoint from Hub and load VAE(base=64).""" logger.info("[load_model] START repo_id=%s filename=%s device=%s", repo_id, filename, DEVICE) path = hf_hub_download(repo_id=repo_id, filename=filename) size_mb = os.path.getsize(path) / (1024 * 1024) if os.path.exists(path) else 0 logger.info("[load_model] checkpoint path=%s size=%.1f MB", path, size_mb) model = VAE(base=64).to(DEVICE) ckpt = torch.load(path, map_location=DEVICE) if isinstance(ckpt, dict) and "model_state_dict" in ckpt: model.load_state_dict(ckpt["model_state_dict"]) logger.info("[load_model] load_state_dict from checkpoint['model_state_dict'] OK") else: model.load_state_dict(ckpt) logger.info("[load_model] load_state_dict from raw state_dict OK") model.eval() nparams = sum(p.numel() for p in model.parameters()) logger.info("[load_model] model.eval() set | params=%s | MODEL LOADED SUCCESSFULLY", nparams) return model def get_model() -> nn.Module: """Return cached model or load once.""" global _model if _model is None: logger.info("[get_model] loading model (first request)") _model = load_model() return _model def preprocess_mask(mask: Image.Image, size: int = INPUT_SIZE) -> torch.Tensor: """ PIL mask -> (1, 4, size, size) float32 in [0, 1]. Steps: Grayscale, Resize(size, size), ToTensor, duplicate to 4 channels, add batch dim. """ logger.info("[preprocess] INPUT size=%s mode=%s", mask.size, mask.mode) transform = T.Compose([ T.Grayscale(num_output_channels=1), T.Resize((size, size), antialias=True), T.ToTensor(), ]) x = transform(mask) # (1, H, W) x_np = x.numpy() logger.info("[preprocess] after Grayscale+Resize(%s)+ToTensor shape=%s min=%.4f max=%.4f mean=%.4f", (size, size), tuple(x.shape), float(x_np.min()), float(x_np.max()), float(x_np.mean())) x = x.repeat(4, 1, 1) # (4, H, W) - duplicate 1 channel to 4 slices x = x.unsqueeze(0) # (1, 4, H, W) - add batch dimension logger.info("[preprocess] after duplicate+batch shape=%s", tuple(x.shape)) return x def mask_to_embedding(mask: Image.Image, size: int = INPUT_SIZE) -> torch.Tensor: """Mask -> preprocess -> encode -> z. Uses train() for forward so BatchNorm uses batch stats (batch size 1). Returns z = mu (no sampling noise) for most faithful reconstruction.""" model = get_model() x = preprocess_mask(mask, size=size).to(DEVICE) model.train() try: with torch.no_grad(): mu, logvar = model.encode(x) z = mu # use mean only for deterministic, most faithful reconstruction (no std*eps) finally: model.eval() return z def decode_to_slices(z: torch.Tensor) -> torch.Tensor: """z -> decode -> (1, 4, 256, 256) batched. Uses train() for forward so BatchNorm uses batch stats (batch size 1).""" model = get_model() model.train() try: with torch.no_grad(): out = model.decode(z) finally: model.eval() out_np = out.detach().cpu().numpy() logger.info("[model output] decode(z) shape=%s min=%.4f max=%.4f mean=%.4f", tuple(out.shape), float(out_np.min()), float(out_np.max()), float(out_np.mean())) return out def enhance_slice(slice_2d: np.ndarray, contrast: bool = True, sharpen: bool = True) -> np.ndarray: """ Postprocess VAE output slice to reduce blur. Input: float32 image [0,1] Output: float32 image [0,1] """ img = slice_2d.astype(np.float32) # --- 1. Contrast stretch (great for CT-like images) --- if contrast: p2, p98 = np.percentile(img, (2, 98)) if p98 > p2: img = (img - p2) / (p98 - p2) img = np.clip(img, 0, 1) # --- 2. Unsharp mask (edge boost) --- if sharpen: blur = cv2.GaussianBlur(img, (0, 0), sigmaX=1.2) img = cv2.addWeighted(img, 1.6, blur, -0.6, 0) return np.clip(img, 0, 1) def inference(mask: Image.Image, slice_index: int = OUTPUT_SLICE_INDEX) -> Tuple[np.ndarray, torch.Tensor]: """ Full pipeline: mask -> encode -> decode -> 4 slices. Returns (one_slice_2d, full_output_tensor). """ z = mask_to_embedding(mask) out = decode_to_slices(z) # (1, 4, 256, 256) batched slice_idx = min(max(0, slice_index), 3) one_slice = out[0, slice_idx].detach().cpu().numpy() # (256, 256) float [0,1] logger.info("[output slice] slice_index=%s shape=%s min=%.4f max=%.4f", slice_idx, one_slice.shape, float(one_slice.min()), float(one_slice.max())) return one_slice, out def slice_to_png(slice_2d: np.ndarray) -> bytes: """(H, W) float [0,1] -> clip -> scale to uint8 -> PNG bytes.""" arr = (np.clip(slice_2d, 0.0, 1.0) * 255).astype(np.uint8) img = Image.fromarray(arr, mode="L") buf = io.BytesIO() img.save(buf, format="PNG") return buf.getvalue() def inference_to_png(mask: Image.Image, slice_index: int = OUTPUT_SLICE_INDEX, contrast: bool = True, sharpen: bool = True) -> bytes: """Mask -> inference -> (optional) enhance -> PNG bytes.""" one_slice, _ = inference(mask, slice_index=slice_index) one_slice = enhance_slice(one_slice, contrast=contrast, sharpen=sharpen) return slice_to_png(one_slice)