segmentation_utils.py

"""Pure image-processing helpers for the street-scene Space."""

from __future__ import annotations

import csv
from pathlib import Path
from typing import Iterable

import numpy as np
from PIL import Image, ImageOps


MAX_OUTPUT_SIDE = 2048

# Official Cityscapes train-ID palette.
CITYSCAPES_PALETTE: dict[str, tuple[int, int, int]] = {
    "road": (128, 64, 128),
    "sidewalk": (244, 35, 232),
    "building": (70, 70, 70),
    "wall": (102, 102, 156),
    "fence": (190, 153, 153),
    "pole": (153, 153, 153),
    "traffic light": (250, 170, 30),
    "traffic sign": (220, 220, 0),
    "vegetation": (107, 142, 35),
    "terrain": (152, 251, 152),
    "sky": (70, 130, 180),
    "person": (220, 20, 60),
    "rider": (255, 0, 0),
    "car": (0, 0, 142),
    "truck": (0, 0, 70),
    "bus": (0, 60, 100),
    "train": (0, 80, 100),
    "motorcycle": (0, 0, 230),
    "bicycle": (119, 11, 32),
}


def _normalise_label(label: str) -> str:
    return label.lower().replace("_", " ").strip()


def _fallback_color(class_id: int) -> tuple[int, int, int]:
    """Return a deterministic, visually distinct color for an unknown class."""
    return (
        int((37 * class_id + 71) % 205 + 25),
        int((67 * class_id + 29) % 205 + 25),
        int((97 * class_id + 11) % 205 + 25),
    )


def class_color(class_id: int, label: str) -> tuple[int, int, int]:
    return CITYSCAPES_PALETTE.get(_normalise_label(label), _fallback_color(class_id))


def resize_for_output(image: Image.Image, max_side: int = MAX_OUTPUT_SIDE) -> Image.Image:
    """Bound output resolution so large phone photos do not exhaust Space memory."""
    image = ImageOps.exif_transpose(image).convert("RGB")
    width, height = image.size
    longest = max(width, height)
    if longest <= max_side:
        return image

    scale = max_side / longest
    size = (max(1, round(width * scale)), max(1, round(height * scale)))
    resampling = getattr(Image, "Resampling", Image)
    return image.resize(size, resampling.LANCZOS)


def render_segmentation(
    image: Image.Image,
    class_map: np.ndarray,
    id2label: dict[int, str],
    opacity: float,
) -> tuple[Image.Image, Image.Image]:
    """Create a Cityscapes-color mask and an overlay with white boundaries."""
    height, width = class_map.shape
    color_array = np.zeros((height, width, 3), dtype=np.uint8)

    for class_id in np.unique(class_map):
        label = id2label.get(int(class_id), f"class_{int(class_id)}")
        color_array[class_map == class_id] = class_color(int(class_id), label)

    base_array = np.asarray(image, dtype=np.float32)
    overlay_array = (
        base_array * (1.0 - opacity) + color_array.astype(np.float32) * opacity
    ).astype(np.uint8)

    boundaries = np.zeros((height, width), dtype=bool)
    boundaries[1:, :] |= class_map[1:, :] != class_map[:-1, :]
    boundaries[:, 1:] |= class_map[:, 1:] != class_map[:, :-1]
    overlay_array[boundaries] = (255, 255, 255)

    return Image.fromarray(overlay_array), Image.fromarray(color_array)


def build_class_table(
    class_map: np.ndarray,
    id2label: dict[int, str],
    min_share_percent: float,
) -> list[list[object]]:
    """Summarise class coverage in descending order."""
    class_ids, counts = np.unique(class_map, return_counts=True)
    total_pixels = int(class_map.size)
    rows: list[list[object]] = []

    for class_id, count in zip(class_ids, counts):
        share = 100.0 * int(count) / total_pixels
        if share < min_share_percent:
            continue
        label = id2label.get(int(class_id), f"class_{int(class_id)}")
        color = class_color(int(class_id), label)
        rows.append(
            [
                int(class_id),
                label,
                int(count),
                round(share, 2),
                "#{:02X}{:02X}{:02X}".format(*color),
            ]
        )

    rows.sort(key=lambda row: float(row[3]), reverse=True)
    return rows


def write_class_csv(path: Path, rows: Iterable[Iterable[object]]) -> None:
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        writer.writerow(["class_id", "class_name", "pixels", "share_percent", "color"])
        writer.writerows(rows)