model.py

# model.py

from __future__ import annotations
import json
from dataclasses import dataclass, asdict
from pathlib import Path
from typing import Dict, List, Tuple, Optional

import numpy as np
from PIL import Image
import nibabel as nib

import torch
import torch.nn as nn
import torch.nn.functional as F
import open_clip  # pip install open_clip_torch


# -----------------------------
# Constants (match your training)
# -----------------------------
REPO_ID = "hf-hub:microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224"

# You trained with "Dementia" as class-3 name (not "AD")
LABEL2IDX: Dict[str, int] = {"CN": 0, "MCI": 1, "Dementia": 2}
IDX2LABEL: Dict[int, str] = {v: k for k, v in LABEL2IDX.items()}


# -----------------------------
# Small config to save with model
# -----------------------------
@dataclass
class ModelConfig:
    model_id: str = REPO_ID
    num_classes: int = 3
    proj_dim: int = 512
    freeze_encoders: bool = False
    label2idx: Dict[str, int] = None

    def to_json(self) -> str:
        d = asdict(self)
        return json.dumps(d, indent=2)

    @staticmethod
    def from_json(path: str | Path) -> "ModelConfig":
        data = json.loads(Path(path).read_text())
        return ModelConfig(**data)


# -----------------------------
# 3D→2D Triptych utilities
# -----------------------------
def center_crop_or_pad(vol: np.ndarray, target_shape: Tuple[int, int, int]) -> np.ndarray:
    """Center-crop or zero-pad a 3D volume to target_shape=(D,H,W)."""
    D, H, W = vol.shape
    tD, tH, tW = target_shape
    out = np.zeros(target_shape, dtype=vol.dtype)

    d0 = max(0, (D - tD) // 2); d1 = d0 + min(D, tD)
    h0 = max(0, (H - tH) // 2); h1 = h0 + min(H, tH)
    w0 = max(0, (W - tW) // 2); w1 = w0 + min(W, tW)

    td0 = max(0, (tD - D) // 2); td1 = td0 + (d1 - d0)
    th0 = max(0, (tH - H) // 2); th1 = th0 + (h1 - h0)
    tw0 = max(0, (tW - W) // 2); tw1 = tw0 + (w1 - w0)

    out[td0:td1, th0:th1, tw0:tw1] = vol[d0:d1, h0:h1, w0:w1]
    return out


def volume_to_triptych(volume_1d: torch.Tensor, out_size: int = 224) -> Image.Image:
    """
    volume_1d: torch tensor [1, D, H, W] in [0,1].
    Returns a PIL RGB image (triptych of axial/coronal/sagittal mid-slices).
    """
    assert volume_1d.ndim == 4 and volume_1d.shape[0] == 1
    _, D, H, W = volume_1d.shape
    v = volume_1d[0].cpu().numpy()  # [D,H,W]

    d_mid, h_mid, w_mid = D // 2, H // 2, W // 2
    axial   = v[d_mid, :, :]        # [H,W]
    coronal = v[:, h_mid, :]        # [D,W] -> resize to [H,W]
    sagitt  = v[:, :, w_mid]        # [D,H] -> resize to [H,W]

    def norm_to_uint8(x: np.ndarray) -> np.ndarray:
        x = (x - x.min()) / (x.max() - x.min() + 1e-8)
        return (x * 255.0).astype(np.uint8)

    axial_img   = Image.fromarray(norm_to_uint8(axial))
    coronal_img = Image.fromarray(norm_to_uint8(coronal)).resize((W, H), Image.BILINEAR)
    sagitt_img  = Image.fromarray(norm_to_uint8(sagitt)).resize((W, H), Image.BILINEAR)

    rgb = np.stack([np.array(axial_img), np.array(coronal_img), np.array(sagitt_img)], axis=-1)
    pil = Image.fromarray(rgb.astype(np.uint8)).resize((out_size, out_size), Image.BILINEAR)
    return pil


# -----------------------------
# The model (same as training)
# -----------------------------
class BiomedClipClassifier(nn.Module):
    """
    Encodes MRI triptych (image) + clinical text with BiomedCLIP (open_clip),
    concatenates L2-normalized embeddings, then classifies into 3 classes.
    """
    def __init__(
        self,
        model_id: str = REPO_ID,
        num_classes: int = 3,
        proj_dim: int = 512,
        freeze_encoders: bool = False,
        device: str = "cpu",
    ):
        super().__init__()
        # Load CLIP model & transforms
        self.clip, self.preprocess_train, self.preprocess_val = open_clip.create_model_and_transforms(model_id)
        self.tokenizer_fn = open_clip.get_tokenizer(model_id)
        self.clip.to(device)

        if freeze_encoders:
            for p in self.clip.parameters():
                p.requires_grad = False

        # Infer feature dims
        with torch.no_grad():
            dummy_img = torch.zeros(1, 3, 224, 224, device=device)
            dummy_txt = self.tokenizer_fn(["test"]).to(device)
            dim_i = self.clip.encode_image(dummy_img).shape[-1]
            dim_t = self.clip.encode_text(dummy_txt).shape[-1]

        in_dim = dim_i + dim_t
        self.head = nn.Sequential(
            nn.Linear(in_dim, proj_dim),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(proj_dim, num_classes),
        )

    def forward(self, images: torch.Tensor, texts_tok: torch.Tensor) -> torch.Tensor:
        img_f = F.normalize(self.clip.encode_image(images), dim=-1)
        txt_f = F.normalize(self.clip.encode_text(texts_tok), dim=-1)
        return self.head(torch.cat([img_f, txt_f], dim=-1))

    # ------------- HF-style save/load -------------
    def save_pretrained(self, save_directory: str | Path, config: Optional[ModelConfig] = None):
        save_dir = Path(save_directory)
        save_dir.mkdir(parents=True, exist_ok=True)
        # state dict
        torch.save(self.state_dict(), save_dir / "pytorch_model.bin")
        # minimal config
        if config is None:
            config = ModelConfig(label2idx=LABEL2IDX)
        (save_dir / "config.json").write_text(config.to_json())

    @staticmethod
    def from_pretrained(load_directory: str | Path, device: str = "cpu") -> "BiomedClipClassifier":
        load_dir = Path(load_directory)
        cfg_path = load_dir / "config.json"
        state_path = load_dir / "pytorch_model.bin"

        if cfg_path.exists():
            cfg = ModelConfig.from_json(cfg_path)
        else:
            # fallback if only a state dict is present
            cfg = ModelConfig(label2idx=LABEL2IDX)

        model = BiomedClipClassifier(
            model_id=cfg.model_id,
            num_classes=cfg.num_classes,
            proj_dim=cfg.proj_dim,
            freeze_encoders=cfg.freeze_encoders,
            device=device,
        )
        if state_path.exists():
            state = torch.load(state_path, map_location=device)
            model.load_state_dict(state, strict=False)
        else:
            # Also allow people to pass a raw .pt file path as directory
            # e.g., repo contains 'biomedclip_best.pt'
            pt_fallback = next(load_dir.glob("*.pt"), None)
            if pt_fallback is not None:
                state = torch.load(pt_fallback, map_location=device)
                model.load_state_dict(state, strict=False)

        model.eval()
        return model


# -----------------------------
# Simple single-sample inference helpers
# -----------------------------
@torch.no_grad()
def predict_from_paths(
    model: BiomedClipClassifier,
    mri_path: str | Path,
    text: str,
    device: str = "cpu",
    use_val_preprocess: bool = True,
    target_shape: Tuple[int, int, int] = (128, 128, 128),
) -> Tuple[str, List[float]]:
    """
    Convenience function to run inference on one NIfTI + text string.
    Returns (pred_label, class_probs).
    """
    model.eval()
    mri_path = Path(mri_path)

    # Load & normalize volume
    vol = nib.load(str(mri_path)).get_fdata().astype(np.float32)
    v = (vol - vol.mean()) / (vol.std() + 1e-8)
    v = (v - v.min()) / (v.max() - v.min() + 1e-8)
    v = center_crop_or_pad(v, target_shape)

    # Triptych -> preprocess
    img_t = torch.from_numpy(v).unsqueeze(0)  # [1,D,H,W]
    trip_pil = volume_to_triptych(img_t)      # PIL RGB 224x224
    preprocess = model.preprocess_val if use_val_preprocess else model.preprocess_train
    img_clip = preprocess(trip_pil).unsqueeze(0).to(device)

    # Tokenize text
    tokenizer = model.tokenizer_fn
    txt_tok = tokenizer([text]).to(device)

    # Forward
    logits = model(img_clip, txt_tok)
    probs = torch.softmax(logits, dim=-1)[0].cpu().tolist()
    pred_idx = int(torch.argmax(logits, dim=-1).item())
    pred_label = IDX2LABEL[pred_idx]
    return pred_label, probs


# -----------------------------
# Minimal example (optional)
# -----------------------------
if __name__ == "__main__":
    # Example: load a local folder with 'pytorch_model.bin' (or a .pt) and run one inference.
    # Set paths before running.
    weights_dir = "./"          # folder containing pytorch_model.bin or a *.pt
    nifti_path  = "/path/to/sample_brain.nii.gz"
    text_input  = "Patient shows mild memory impairment and hippocampal atrophy."

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = BiomedClipClassifier.from_pretrained(weights_dir, device=device)

    pred, probs = predict_from_paths(model, nifti_path, text_input, device=device)
    print("Prediction:", pred)
    print("Probabilities [CN, MCI, Dementia]:", [round(p, 4) for p in probs])