inference_core.py

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import albumentations as A
import cv2
import numpy as np
import rasterio.features
import torch
import torch.nn as nn
import yaml
from albumentations.pytorch import ToTensorV2
from dynamic_network_architectures.architectures.unet import PlainConvUNet
from huggingface_hub import hf_hub_download
from matplotlib import colormaps
from shapely.geometry import shape
from shapely.ops import unary_union

MODEL_REPO_ID = "AImageLab-Zip/CALHippo-Framework-Models"
MODEL_CONFIG_PATH = (
    "density_estimation/short_unet/"
    "9_shorter_unet_normalizedgame_asymclassnormalizedl1loss_adamw.yaml"
)
MODEL_WEIGHTS_PATH = "density_estimation/short_unet/final_density_model.pth"

CLASS_NAMES = ["Pyramidal", "Interneuron", "Astrocyte"]
CLASS_COLORS = np.array(
    [
        (214, 39, 40),
        (0, 153, 170),
        (31, 119, 180),
    ],
    dtype=np.uint8,
)


@dataclass(frozen=True)
class LoadedModel:
    model: nn.Module
    transform: A.Compose
    max_pix_value: float
    patch_size: int
    stride: int
    num_classes: int
    device: str


class PlainConvUNetReLU(nn.Module):
    def __init__(
        self,
        base: PlainConvUNet,
        num_classes: int,
        output_scalers: list[float] | None = None,
        output_activation: str = "relu",
    ) -> None:
        super().__init__()
        self.base = base
        self.output_scaler = None
        if output_scalers is not None:
            self.output_scaler = nn.Parameter(
                torch.tensor(output_scalers, dtype=torch.float32)
            )

        if output_activation.lower() == "relu":
            self.output_act = nn.ReLU(inplace=False)
        elif output_activation.lower() == "softplus":
            self.output_act = nn.Softplus()
        elif output_activation.lower() == "none":
            self.output_act = nn.Identity()
        else:
            raise ValueError(f"Unsupported output activation: {output_activation}")

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = self.base(x)
        if self.output_scaler is not None:
            out = out * self.output_scaler.view(1, -1, 1, 1)
        return self.output_act(out)


def _build_model(config: dict[str, Any]) -> nn.Module:
    model_config = config.get("MODEL", {})
    kwargs = dict(model_config.get("kwargs", {}))
    output_scalers = kwargs.pop("output_scalers", None)
    output_activation = kwargs.pop("output_activation", "relu")
    kwargs.pop("use_log_counts", None)

    norm_ops = {
        "BatchNorm2d": nn.BatchNorm2d,
        "InstanceNorm2d": nn.InstanceNorm2d,
    }
    nonlins = {
        "LeakyReLU": nn.LeakyReLU,
        "ReLU": nn.ReLU,
        "GELU": nn.GELU,
        "PReLU": nn.PReLU,
        "ELU": nn.ELU,
        "SiLU": nn.SiLU,
    }

    kwargs["conv_op"] = nn.Conv2d
    kwargs["dropout_op"] = None
    kwargs["norm_op"] = norm_ops[kwargs.get("norm_op", "InstanceNorm2d")]
    kwargs["nonlin"] = nonlins[kwargs.get("nonlin", "LeakyReLU")]

    num_classes = int(model_config.get("num_classes", 3))
    base = PlainConvUNet(
        input_channels=int(model_config.get("input_channels", 3)),
        num_classes=num_classes,
        deep_supervision=bool(model_config.get("deep_supervision", False)),
        **kwargs,
    )
    return PlainConvUNetReLU(
        base=base,
        num_classes=num_classes,
        output_scalers=output_scalers,
        output_activation=output_activation,
    )


def load_demo_model() -> LoadedModel:
    config_path = hf_hub_download(MODEL_REPO_ID, MODEL_CONFIG_PATH)
    weights_path = hf_hub_download(MODEL_REPO_ID, MODEL_WEIGHTS_PATH)

    with Path(config_path).open("r") as fh:
        config = yaml.safe_load(fh) or {}

    data_config = config.get("DATA", {})
    patch_size = int(data_config.get("img_size", 128))
    max_pix_value = float(data_config.get("fill_value", 65535))
    norm_mean = tuple(data_config.get("norm_mean", [0.7637, 0.7637, 0.7637]))
    norm_std = tuple(data_config.get("norm_std", [0.0703, 0.0703, 0.0703]))
    model_config = config.get("MODEL", {})
    num_classes = int(model_config.get("num_classes", 3))

    transform = A.Compose(
        [
            A.PadIfNeeded(
                min_height=patch_size,
                min_width=patch_size,
                border_mode=cv2.BORDER_CONSTANT,
                fill=1.0,
                fill_mask=0,
            ),
            A.Normalize(mean=norm_mean, std=norm_std, max_pixel_value=1.0),
            ToTensorV2(transpose_mask=True),
        ]
    )

    # The public demo targets the free CPU Basic Space tier.
    device = "cpu"
    model = _build_model(config).to(device)
    try:
        state_dict = torch.load(weights_path, map_location=device, weights_only=True)
    except TypeError:
        state_dict = torch.load(weights_path, map_location=device)
    model.load_state_dict(state_dict)
    model.eval()

    return LoadedModel(
        model=model,
        transform=transform,
        max_pix_value=max_pix_value,
        patch_size=patch_size,
        stride=patch_size // 2,
        num_classes=num_classes,
        device=device,
    )


def load_low_res_wsi(
    image_path: str | Path,
    max_pix_value: float,
    transform: A.Compose | None = None,
) -> tuple[np.ndarray | torch.Tensor, dict[str, int]]:
    wsi = cv2.imread(str(image_path), cv2.IMREAD_UNCHANGED)
    if wsi is None:
        raise ValueError(f"Could not read image: {image_path}")
    if wsi.ndim == 2:
        wsi = cv2.cvtColor(wsi, cv2.COLOR_GRAY2RGB)
    else:
        wsi = cv2.cvtColor(wsi, cv2.COLOR_BGR2RGB)
    wsi = wsi.astype(np.float32) / max_pix_value

    pad = {"top": 0, "bottom": 0, "left": 0, "right": 0}
    if transform is None:
        return wsi, pad

    h_orig, w_orig = wsi.shape[:2]
    augmented = transform(image=wsi)
    tensor = augmented["image"]
    h_new, w_new = tensor.shape[1], tensor.shape[2]

    if h_new > h_orig:
        diff_h = h_new - h_orig
        pad["top"] = diff_h // 2
        pad["bottom"] = diff_h - pad["top"]
    if w_new > w_orig:
        diff_w = w_new - w_orig
        pad["left"] = diff_w // 2
        pad["right"] = diff_w - pad["left"]

    return tensor, pad


def create_gaussian_mask(
    patch_size: int,
    sigma: float,
    device: str,
) -> torch.Tensor:
    coords = torch.arange(patch_size, dtype=torch.float32, device=device)
    coords -= (patch_size - 1) / 2.0
    y, x = torch.meshgrid(coords, coords, indexing="ij")
    return torch.exp(-(x**2 + y**2) / (2 * sigma**2))


def predict_density_map(
    wsi_tensor: torch.Tensor,
    loaded: LoadedModel,
    inference_batch_size: int = 8,
) -> torch.Tensor:
    _, _, height, width = wsi_tensor.shape
    patch_size = loaded.patch_size
    stride = loaded.stride
    device = loaded.device

    global_density = torch.zeros(
        (1, loaded.num_classes, height, width), dtype=torch.float32, device=device
    )
    global_weight = torch.zeros_like(global_density)
    gaussian = create_gaussian_mask(
        patch_size=patch_size,
        sigma=(patch_size // 2) // 3,
        device=device,
    ).view(1, 1, patch_size, patch_size)

    max_h = height - patch_size
    max_w = width - patch_size
    anchors = [
        (y, x)
        for y in list(range(0, max_h, stride)) + [max_h]
        for x in list(range(0, max_w, stride)) + [max_w]
    ]
    patches = [
        wsi_tensor[:, :, y : y + patch_size, x : x + patch_size] for y, x in anchors
    ]

    with torch.no_grad():
        for start in range(0, len(patches), inference_batch_size):
            batch = torch.cat(patches[start : start + inference_batch_size], dim=0).to(
                device
            )
            preds = loaded.model(batch)
            batch_anchors = anchors[start : start + inference_batch_size]
            for pred, (y, x) in zip(preds, batch_anchors):
                pred = pred.unsqueeze(0)
                global_density[:, :, y : y + patch_size, x : x + patch_size] += (
                    pred * gaussian
                )
                global_weight[:, :, y : y + patch_size, x : x + patch_size] += gaussian

    return global_density / (global_weight + 1e-8)


def unpad_density_map(
    density: torch.Tensor,
    pad: dict[str, int],
) -> torch.Tensor:
    y_end = density.shape[2] - pad["bottom"]
    x_end = density.shape[3] - pad["right"]
    return density[:, :, pad["top"] : y_end, pad["left"] : x_end]


def extract_roi_mask_from_geojson(
    geojson_path: str | Path,
    image_shape: tuple[int, int],
    roi_class: str = "OverallCA",
) -> np.ndarray:
    with Path(geojson_path).open("r") as fh:
        geojson = json.load(fh)

    roi_geoms = []
    for feature in geojson.get("features", []):
        props = feature.get("properties", {})
        object_class = props.get("classification", {}).get("name")
        if object_class != roi_class:
            continue
        geom = shape(feature.get("geometry", {}))
        if geom.is_valid:
            roi_geoms.append(geom)

    if not roi_geoms:
        return np.zeros(image_shape, dtype=np.uint8)

    return rasterio.features.rasterize([unary_union(roi_geoms)], out_shape=image_shape)


def sample_discrete_density_numpy(
    density_mask: np.ndarray,
    rng: np.random.Generator,
) -> np.ndarray:
    height, width, channels = density_mask.shape
    discrete = np.zeros((height, width, channels), dtype=np.int32)
    clean_density = np.clip(density_mask, a_min=0.0, a_max=None)

    for channel_idx in range(channels):
        channel_map = clean_density[:, :, channel_idx]
        total_mass = float(channel_map.sum())
        if total_mass <= 1e-6:
            continue

        int_mass = int(total_mass)
        extra_cell = 1 if rng.random() < total_mass - int_mass else 0
        n_samples = int_mass + extra_cell
        if n_samples == 0:
            continue

        pmf = (channel_map / total_mass).ravel()
        pmf[-1] = 1.0 - pmf[:-1].sum()
        pmf = np.clip(pmf, a_min=0.0, a_max=None)
        pmf = pmf / pmf.sum()
        discrete[:, :, channel_idx] = rng.multinomial(n_samples, pmf).reshape(
            height, width
        )

    return discrete


def normalize_original_for_display(image_path: str | Path) -> np.ndarray:
    image = cv2.imread(str(image_path), cv2.IMREAD_UNCHANGED)
    if image.ndim == 2:
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
    else:
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    image = image.astype(np.float32)
    image -= image.min()
    max_value = image.max()
    if max_value > 0:
        image /= max_value
    return (image * 255).astype(np.uint8)


def colorize_density(density: np.ndarray) -> np.ndarray:
    density = np.clip(density, 0, None)
    if density.max() > density.min():
        scaled = (density - density.min()) / (density.max() - density.min())
    else:
        scaled = np.zeros_like(density)
    return (colormaps["viridis"](scaled)[..., :3] * 255).astype(np.uint8)


def normalize_density_channel(channel: np.ndarray) -> np.ndarray:
    channel = np.clip(channel.astype(np.float32), 0, None)
    positive = channel[channel > 0]
    if positive.size == 0:
        return np.zeros_like(channel, dtype=np.float32)

    lo = float(np.percentile(positive, 2))
    hi = float(np.percentile(positive, 99))
    if hi <= lo:
        hi = float(np.max(positive))
        lo = 0.0
    if hi <= lo:
        return np.clip(channel, 0.0, 1.0)

    normalized = (channel - lo) / (hi - lo)
    normalized[channel <= 0] = 0.0
    return np.clip(normalized, 0.0, 1.0)


def colorize_sampled(sampled: np.ndarray, color: np.ndarray) -> np.ndarray:
    image = np.full((*sampled.shape, 3), 255, dtype=np.uint8)
    mask = sampled > 0
    image[mask] = color
    return image


def build_combined_points(sampled: np.ndarray) -> np.ndarray:
    combined = np.full((*sampled.shape[:2], 3), 255, dtype=np.uint8)
    for class_idx, color in enumerate(CLASS_COLORS):
        combined[sampled[:, :, class_idx] > 0] = color
    return combined


def build_combined_density_overlay(
    image: np.ndarray,
    density: np.ndarray,
    alpha: float = 0.72,
) -> np.ndarray:
    overlay = np.zeros((*density.shape[:2], 3), dtype=np.float32)
    for class_idx, color in enumerate(CLASS_COLORS):
        normalized = normalize_density_channel(density[:, :, class_idx])
        color_arr = color.astype(np.float32) / 255.0
        overlay += normalized[..., np.newaxis] * color_arr

    overlay = np.clip(overlay, 0.0, 1.0)
    base = image.astype(np.float32) / 255.0
    blended = base * (1.0 - alpha) + overlay * alpha
    blended = np.clip(blended, 0.0, 1.0)
    return np.rint(blended * 255.0).astype(np.uint8)


def run_demo_inference(
    image_path: str | Path,
    geojson_path: str | Path | None = None,
    use_roi: bool = True,
    seed: int = 42,
) -> dict[str, Any]:
    loaded = load_demo_model()
    wsi_tensor, pad = load_low_res_wsi(
        image_path=image_path,
        max_pix_value=loaded.max_pix_value,
        transform=loaded.transform,
    )
    wsi_tensor = wsi_tensor.unsqueeze(0).to(loaded.device)
    density = predict_density_map(wsi_tensor, loaded=loaded)
    density = unpad_density_map(density, pad)
    density_array = density.squeeze(0).permute(1, 2, 0).cpu().numpy()

    roi_mask = None
    pred_array = density_array
    if use_roi and geojson_path is not None:
        roi_mask = extract_roi_mask_from_geojson(
            geojson_path=geojson_path,
            image_shape=density_array.shape[:2],
        )
        pred_array = density_array * roi_mask[..., np.newaxis]

    sampled = sample_discrete_density_numpy(pred_array, rng=np.random.default_rng(seed))

    density_maps = [
        colorize_density(pred_array[:, :, i]) for i in range(loaded.num_classes)
    ]
    sampled_maps = [
        colorize_sampled(sampled[:, :, i], CLASS_COLORS[i])
        for i in range(loaded.num_classes)
    ]
    counts = []
    for idx, class_name in enumerate(CLASS_NAMES[: loaded.num_classes]):
        counts.append(
            {
                "class": class_name,
                "density_sum": round(float(pred_array[:, :, idx].sum()), 2),
                "sampled_count": int(sampled[:, :, idx].sum()),
            }
        )

    original = normalize_original_for_display(image_path)

    return {
        "original": original,
        "roi_mask": roi_mask,
        "combined_density": build_combined_density_overlay(original, pred_array),
        "density_maps": density_maps,
        "sampled_maps": sampled_maps,
        "combined_points": build_combined_points(sampled),
        "counts": counts,
    }