inference.py

"""
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)