Spaces:
Runtime error
Runtime error
File size: 1,989 Bytes
e0e2b27 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 | """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
|