scripts/segmentation_utils.py

"""Shared helpers for YOLO-polygon semantic segmentation workflows."""

from __future__ import annotations

from pathlib import Path

import cv2
import numpy as np
import yaml


def load_dataset_class_names(dataset: Path) -> dict[int, str]:
    data_yaml = dataset / "data.yaml"
    data = yaml.safe_load(data_yaml.read_text(encoding="utf-8"))
    return {int(k): str(v) for k, v in data["names"].items()}


def load_data_yaml_class_names(data_yaml: Path) -> dict[int, str]:
    data = yaml.safe_load(data_yaml.read_text(encoding="utf-8"))
    return {int(k): str(v) for k, v in data["names"].items()}


def yolo_label_to_semantic_mask(
    label_path: Path,
    height: int,
    width: int,
    *,
    background_value: int = -1,
    class_offset: int = 0,
) -> np.ndarray:
    mask = np.full((height, width), background_value, dtype=np.int16)
    if not label_path.exists():
        return mask

    for line in label_path.read_text(encoding="utf-8").splitlines():
        parts = line.split()
        if len(parts) < 7:
            continue
        cls = int(parts[0]) + class_offset
        coords = np.array(parts[1:], dtype=np.float32).reshape(-1, 2)
        coords[:, 0] *= width
        coords[:, 1] *= height
        cv2.fillPoly(mask, [coords.astype(np.int32)], cls)
    return mask


def yolo_result_to_semantic_mask(
    result,
    height: int,
    width: int,
    *,
    mask_threshold: float = 0.5,
    background_value: int = -1,
) -> np.ndarray:
    mask = np.full((height, width), background_value, dtype=np.int16)
    if result.masks is None:
        return mask

    masks = result.masks.data.cpu().numpy()
    classes = result.boxes.cls.cpu().numpy().astype(int)
    confs = result.boxes.conf.cpu().numpy()

    for idx in np.argsort(confs):
        item = masks[idx]
        if item.shape != (height, width):
            item = cv2.resize(item, (width, height), interpolation=cv2.INTER_NEAREST)
        mask[item > mask_threshold] = int(classes[idx])
    return mask