utils/qualitative.py

from __future__ import annotations

import json
import random
import re
from pathlib import Path
from typing import Iterable

import numpy as np
import torch
from PIL import Image, ImageDraw

from datasets.cd_dataset import CDDataset


ROOT = Path(__file__).resolve().parents[1]


def safe_sample_id(sample_id: str) -> str:
    return re.sub(r"[^A-Za-z0-9_.-]+", "_", str(sample_id))


def _boundary_density(mask: np.ndarray) -> float:
    mask = mask.astype(bool)
    if not mask.any():
        return 0.0
    horiz = np.zeros_like(mask, dtype=bool)
    vert = np.zeros_like(mask, dtype=bool)
    horiz[:, 1:] = mask[:, 1:] != mask[:, :-1]
    vert[1:, :] = mask[1:, :] != mask[:-1, :]
    return float((horiz | vert).sum()) / float(mask.size)


def qualitative_manifest_path(dataset_name: str) -> Path:
    return ROOT / "results" / "qualitative_samples" / dataset_name / "sample_manifest.json"


def select_or_load_manifest(
    dataset_cfg: dict,
    count: int = 20,
    seed: int = 3407,
    force: bool = False,
) -> dict:
    dataset_name = dataset_cfg["name"]
    path = qualitative_manifest_path(dataset_name)
    if path.exists() and not force:
        with path.open("r", encoding="utf-8") as f:
            return json.load(f)

    ds = CDDataset(dataset_cfg["data_root"], "test", cfg=dataset_cfg, normalize=False, return_format="tuple")
    rng = random.Random(seed)
    rows = []
    for index, (a_path, b_path, mask_path, sample_id) in enumerate(ds.samples):
        if not (a_path.is_file() and b_path.is_file() and mask_path.is_file()):
            continue
        mask = np.asarray(Image.open(mask_path).convert("L").resize((ds.image_size, ds.image_size), Image.NEAREST))
        binary = mask > (0 if mask.max() <= 1 else ds.threshold)
        ratio = float(binary.mean())
        boundary = _boundary_density(binary)
        if ratio <= 0.0 or ratio >= 0.85 or boundary <= 0.0:
            continue
        rows.append({
            "index": index,
            "sample_id": sample_id,
            "a_path": str(a_path),
            "b_path": str(b_path),
            "mask_path": str(mask_path),
            "changed_ratio": ratio,
            "boundary_density": boundary,
            "tie": rng.random(),
        })

    if len(ds) >= count and len(rows) < count:
        raise RuntimeError(
            f"Only {len(rows)} visually useful test samples found for {dataset_name}, "
            f"but {count} are required from {len(ds)} total test samples."
        )

    rows = sorted(rows, key=lambda r: (-r["boundary_density"], -r["changed_ratio"], r["sample_id"], r["tie"]))
    selected = rows[: min(count, len(rows))]
    for rank, row in enumerate(selected, start=1):
        row["rank"] = rank
        row.pop("tie", None)
    path.parent.mkdir(parents=True, exist_ok=True)
    manifest = {
        "dataset": dataset_name,
        "split": "test",
        "seed": seed,
        "requested_count": count,
        "selected_count": len(selected),
        "total_test_samples": len(ds),
        "selection_rule": "non-empty mask, changed_ratio < 0.85, positive boundary density, sorted deterministically",
        "samples": selected,
    }
    with path.open("w", encoding="utf-8") as f:
        json.dump(manifest, f, indent=2, sort_keys=True)
    return manifest


def manifest_ids(manifest: dict) -> set[str]:
    return {str(row["sample_id"]) for row in manifest.get("samples", [])}


def denormalize(tensor: torch.Tensor, mean: list[float], std: list[float]) -> torch.Tensor:
    if tensor.ndim == 3:
        mean_t = torch.tensor(mean, dtype=tensor.dtype, device=tensor.device).view(3, 1, 1)
        std_t = torch.tensor(std, dtype=tensor.dtype, device=tensor.device).view(3, 1, 1)
    else:
        mean_t = torch.tensor(mean, dtype=tensor.dtype, device=tensor.device).view(1, 3, 1, 1)
        std_t = torch.tensor(std, dtype=tensor.dtype, device=tensor.device).view(1, 3, 1, 1)
    return tensor * std_t + mean_t


def tensor_to_rgb_image(tensor: torch.Tensor) -> Image.Image:
    if tensor.ndim == 4:
        tensor = tensor[0]
    arr = tensor.detach().cpu().float().clamp(0, 1).numpy()
    if arr.shape[0] == 1:
        arr = np.repeat(arr, 3, axis=0)
    arr = np.transpose(arr[:3], (1, 2, 0))
    return Image.fromarray((arr * 255).astype(np.uint8), mode="RGB")


def mask_to_image(mask: torch.Tensor | np.ndarray) -> Image.Image:
    if torch.is_tensor(mask):
        arr = mask.detach().cpu().numpy()
    else:
        arr = mask
    arr = np.squeeze(arr)
    arr = (arr > 0).astype(np.uint8) * 255
    return Image.fromarray(arr, mode="L").convert("RGB")


def overlay_image(gt: torch.Tensor, pred: torch.Tensor) -> Image.Image:
    gt_arr = np.squeeze(gt.detach().cpu().numpy()).astype(bool)
    pred_arr = np.squeeze(pred.detach().cpu().numpy()).astype(bool)
    rgb = np.zeros((*gt_arr.shape, 3), dtype=np.uint8)
    rgb[gt_arr & pred_arr] = (255, 255, 255)
    rgb[pred_arr & ~gt_arr] = (255, 64, 64)
    rgb[gt_arr & ~pred_arr] = (64, 160, 255)
    return Image.fromarray(rgb, mode="RGB")


def save_binary_prediction(pred: torch.Tensor, path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    mask_to_image(pred).convert("L").save(path)


def save_probability_map(prob: torch.Tensor, path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    arr = np.squeeze(prob.detach().cpu().float().clamp(0, 1).numpy())
    Image.fromarray((arr * 255).astype(np.uint8), mode="L").save(path)


def save_visual_panel(
    a: torch.Tensor,
    b: torch.Tensor,
    gt: torch.Tensor,
    pred: torch.Tensor,
    out_path: Path,
    prob: torch.Tensor | None = None,
) -> None:
    tiles = [
        ("A", tensor_to_rgb_image(a)),
        ("B", tensor_to_rgb_image(b)),
        ("GT", mask_to_image(gt)),
        ("Pred", mask_to_image(pred)),
        ("Overlay", overlay_image(gt, pred)),
    ]
    if prob is not None:
        tiles.insert(4, ("Score", mask_to_image((prob >= 0.5).float())))
    tile_w, tile_h = tiles[0][1].size
    label_h = 24
    panel = Image.new("RGB", (tile_w * len(tiles), tile_h + label_h), "white")
    draw = ImageDraw.Draw(panel)
    for i, (label, img) in enumerate(tiles):
        x = i * tile_w
        panel.paste(img.resize((tile_w, tile_h), Image.NEAREST), (x, label_h))
        draw.text((x + 6, 5), label, fill=(0, 0, 0))
    out_path.parent.mkdir(parents=True, exist_ok=True)
    panel.save(out_path)


def rank_for_sample(manifest: dict, sample_id: str) -> int | None:
    for row in manifest.get("samples", []):
        if str(row["sample_id"]) == str(sample_id):
            return int(row["rank"])
    return None