model_loader.py

from pathlib import Path
from typing import Any

import numpy as np
from PIL import Image

try:
    import torch
except Exception:  # pragma: no cover - torch is optional for fallback mode
    torch = None

try:
    import cv2
except Exception:  # pragma: no cover - OpenCV is optional
    cv2 = None


MODEL_PATHS = [
    Path(__file__).resolve().parent / "project" / "best_model.pth",
    Path(__file__).resolve().parent / "XPLENDIA_GitHub" / "segmentation_head.pth",
]
INPUT_SIZE = (256, 256)

CLASS_COLORS = np.array([
    [18, 18, 18],      # Background
    [117, 117, 117],   # Road / Track
    [0, 200, 83],      # Grass
    [0, 122, 51],      # Tree
    [0, 120, 215],     # Vehicle
    [231, 76, 60],     # Obstacle
    [66, 165, 245],    # Sky
    [141, 110, 99],    # Building
    [236, 64, 122],    # Person
    [255, 167, 38],    # Other Terrain
], dtype=np.uint8)


def _get_device() -> Any:
    if torch is None:
        return "cpu"
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")


def _as_rgb_array(image: Any) -> np.ndarray:
    if isinstance(image, Image.Image):
        return np.array(image.convert("RGB"), dtype=np.uint8)

    arr = np.asarray(image)
    if arr.ndim == 2:
        arr = np.stack([arr] * 3, axis=-1)
    elif arr.ndim == 3 and arr.shape[2] >= 3:
        arr = arr[:, :, :3]
    else:
        raise ValueError("Unsupported image shape")

    if arr.dtype != np.uint8:
        arr = arr.astype(np.float32)
        if arr.max() <= 1.5:
            arr = arr * 255.0
        arr = np.clip(arr, 0, 255).astype(np.uint8)
    return arr


def preprocess_image(image: Any) -> np.ndarray:
    """Resize and normalize an image to a 256x256 float array in [0, 1]."""
    rgb = _as_rgb_array(image)
    resized = Image.fromarray(rgb).resize(INPUT_SIZE, Image.BILINEAR)
    return np.asarray(resized, dtype=np.float32) / 255.0


def resize_mask_nearest(mask: Any, target_size: tuple[int, int]) -> np.ndarray:
    """Resize a class-id mask with nearest-neighbor interpolation."""
    mask_ids = mask_to_class_ids(mask)
    target_w, target_h = target_size
    if cv2 is not None:
        return cv2.resize(mask_ids.astype(np.uint8), (target_w, target_h), interpolation=cv2.INTER_NEAREST)
    resized = Image.fromarray(mask_ids.astype(np.uint8), mode="L").resize((target_w, target_h), Image.NEAREST)
    return np.asarray(resized, dtype=np.uint8)


class DummyModel:
    """Safe fallback model that produces a deterministic terrain-like mask."""

    def predict(self, image: Any) -> np.ndarray:
        arr = np.asarray(image, dtype=np.float32)
        if arr.ndim != 3:
            raise ValueError("DummyModel expects an RGB image array")

        if arr.max() > 1.5:
            arr = arr / 255.0

        rgb = arr[:, :, :3]
        r = rgb[:, :, 0]
        g = rgb[:, :, 1]
        b = rgb[:, :, 2]
        gray = rgb.mean(axis=2)
        saturation = rgb.max(axis=2) - rgb.min(axis=2)

        h, w = gray.shape
        yy = np.linspace(0.0, 1.0, h, dtype=np.float32)[:, None]
        xx = np.linspace(0.0, 1.0, w, dtype=np.float32)[None, :]

        mask = np.full((h, w), 9, dtype=np.uint8)
        mask[gray < 0.18] = 0

        road = (saturation < 0.18) & (gray >= 0.18) & (gray < 0.78)
        mask[road] = 1

        grass = (g > r + 0.06) & (g > b + 0.03) & (gray >= 0.18) & (gray < 0.8)
        mask[grass] = 2

        tree = grass & (gray < 0.45)
        mask[tree] = 3

        vehicle = (b > r + 0.12) & (b > g + 0.05) & (gray > 0.22)
        mask[vehicle] = 4

        obstacle = (r > g + 0.12) & (r > b + 0.08) & (gray > 0.2)
        mask[obstacle] = 5

        sky = (b > r + 0.08) & (b > g + 0.04) & (gray > 0.5)
        mask[sky] = 6

        building = (saturation < 0.12) & (gray > 0.35) & (gray < 0.68)
        mask[building] = 7

        person = (r > 0.45) & (b > 0.2) & (g < 0.55) & (gray > 0.25)
        mask[person] = 8

        road_band = (yy > 0.35) & (yy < 0.72) & (mask == 9)
        mask[road_band] = 1

        terrain = (mask == 9) & (xx > 0.05) & (xx < 0.95)
        mask[terrain] = 9
        return mask


def load_model() -> Any:
    """Load a real model if possible, otherwise return DummyModel."""
    if torch is not None:
        for model_path in MODEL_PATHS:
            if not model_path.exists():
                continue
            try:
                model = torch.load(model_path, map_location=_get_device())
                if hasattr(model, "eval") and callable(model):
                    model.eval()
                    return model
                if hasattr(model, "predict") or callable(model):
                    return model
            except Exception:
                continue
    return DummyModel()


def _output_to_mask(output: Any) -> np.ndarray:
    if torch is not None and isinstance(output, torch.Tensor):
        tensor = output.detach().cpu()
        if tensor.dim() == 4:
            tensor = torch.argmax(tensor, dim=1).squeeze(0)
        elif tensor.dim() == 3 and tensor.shape[0] > 1:
            tensor = torch.argmax(tensor, dim=0)
        elif tensor.dim() == 3:
            tensor = tensor.squeeze(0).long()
        elif tensor.dim() == 2:
            tensor = tensor.long()
        else:
            raise ValueError(f"Unexpected tensor shape: {tuple(tensor.shape)}")
        return tensor.numpy().astype(np.uint8)

    arr = np.asarray(output)
    if arr.ndim == 4:
        arr = arr[0]
    if arr.ndim == 3 and arr.shape[0] > 1 and arr.shape[0] <= 32:
        arr = np.argmax(arr, axis=0)
    elif arr.ndim == 3 and arr.shape[-1] > 1 and arr.shape[-1] <= 32:
        arr = np.argmax(arr, axis=-1)
    elif arr.ndim == 3:
        arr = arr.squeeze()
    elif arr.ndim != 2:
        raise ValueError(f"Unsupported output shape: {arr.shape}")
    return arr.astype(np.uint8)


def mask_to_class_ids(mask: Any) -> np.ndarray:
    """Convert a grayscale or RGB mask into class-id form."""
    if isinstance(mask, Image.Image):
        arr = np.array(mask)
    else:
        arr = np.asarray(mask)

    if arr.ndim == 2:
        if arr.dtype.kind == "f" and arr.max() <= 1.0:
            arr = np.rint(arr * 9.0)
        unique = np.unique(arr)
        if unique.size <= 2:
            if arr.max() <= 1:
                return np.clip(arr, 0, 9).astype(np.uint8)
            return (arr > 0).astype(np.uint8)
        if arr.max() <= 9:
            return np.clip(arr, 0, 9).astype(np.uint8)
        return np.clip(np.rint(arr.astype(np.float32) / 255.0 * 9.0), 0, 9).astype(np.uint8)

    if arr.ndim == 3:
        rgb = arr[:, :, :3].astype(np.int16)
        palette = CLASS_COLORS.astype(np.int16)
        distances = np.sqrt(((rgb[:, :, None, :] - palette[None, None, :, :]) ** 2).sum(axis=3))
        return np.argmin(distances, axis=2).astype(np.uint8)

    raise ValueError("Mask must be a 2D or 3D array/image")


def mask_to_color_image(mask: Any) -> Image.Image:
    class_ids = mask_to_class_ids(mask)
    rgb = CLASS_COLORS[class_ids]
    return Image.fromarray(rgb.astype(np.uint8), mode="RGB")


def overlay_mask_on_image(image: Any, mask: Any, alpha: float = 0.45) -> Image.Image:
    base = Image.fromarray(_as_rgb_array(image)).convert("RGB")
    resized_mask = resize_mask_nearest(mask, base.size)
    color_mask = np.asarray(mask_to_color_image(resized_mask), dtype=np.uint8)
    base_arr = np.asarray(base, dtype=np.uint8)

    if cv2 is not None:
        overlay_bgr = cv2.addWeighted(
            cv2.cvtColor(base_arr, cv2.COLOR_RGB2BGR),
            1.0 - alpha,
            cv2.cvtColor(color_mask, cv2.COLOR_RGB2BGR),
            alpha,
            0,
        )
        return Image.fromarray(cv2.cvtColor(overlay_bgr, cv2.COLOR_BGR2RGB))

    overlay = (base_arr.astype(np.float32) * (1.0 - alpha) + color_mask.astype(np.float32) * alpha).astype(np.uint8)
    return Image.fromarray(overlay)


def calculate_binary_iou(pred_mask: Any, true_mask: Any) -> float:
    """Calculate IoU for binary masks."""
    pred_ids = (mask_to_class_ids(pred_mask) > 0).astype(np.uint8)
    true_ids = (mask_to_class_ids(true_mask) > 0).astype(np.uint8)

    intersection = np.logical_and(pred_ids, true_ids).sum()
    union = np.logical_or(pred_ids, true_ids).sum()
    if union == 0:
        return 1.0
    return float(intersection / union)


def calculate_mean_iou(pred_mask: Any, true_mask: Any, num_classes: int = 10) -> float:
    """Calculate mean IoU across all classes."""
    pred_ids = mask_to_class_ids(pred_mask)
    true_ids = mask_to_class_ids(true_mask)

    if pred_ids.shape != true_ids.shape:
        resized = Image.fromarray(true_ids.astype(np.uint8)).resize(pred_ids.shape[::-1], Image.NEAREST)
        true_ids = np.asarray(resized, dtype=np.uint8)

    ious = []
    for class_id in range(num_classes):
        pred_class = pred_ids == class_id
        true_class = true_ids == class_id
        union = np.logical_or(pred_class, true_class).sum()
        if union == 0:
            continue
        intersection = np.logical_and(pred_class, true_class).sum()
        ious.append(intersection / union)

    if not ious:
        return 1.0
    return float(np.mean(ious))


def predict(model: Any, image: Any) -> np.ndarray:
    """Run segmentation and return a 2D class-id mask."""
    prepared = preprocess_image(image)

    if hasattr(model, "predict"):
        output = model.predict(prepared)
        return _output_to_mask(output)

    if callable(model):
        if torch is not None:
            tensor = torch.from_numpy(prepared).permute(2, 0, 1).unsqueeze(0).float().to(_get_device())
            try:
                with torch.no_grad():
                    output = model(tensor)
            except Exception:
                output = model(prepared)
        else:
            output = model(prepared)
        if isinstance(output, (tuple, list)):
            output = output[0]
        if isinstance(output, dict):
            output = next(iter(output.values()))
        return _output_to_mask(output)

    raise TypeError(f"Unsupported model object: {type(model)}")


def calculate_iou(pred_mask: Any, true_mask: Any | None = None, num_classes: int = 10) -> float | None:
    """Compute IoU for the current segmentation setup."""
    if true_mask is None:
        return None

    pred_ids = mask_to_class_ids(pred_mask)
    true_ids = mask_to_class_ids(true_mask)

    is_binary = np.unique(pred_ids).size <= 2 and np.unique(true_ids).size <= 2
    if is_binary:
        return calculate_binary_iou(pred_ids, true_ids)

    return calculate_mean_iou(pred_ids, true_ids, num_classes=num_classes)