82 lines
3.8 KiB
Python
82 lines
3.8 KiB
Python
"""Supervised classification from fixed-size YOLO tiles and human box labels."""
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from pathlib import Path
|
|
|
|
import yaml
|
|
from PIL import Image
|
|
|
|
from datasets.classification import FolderClassificationDataset
|
|
from datasets.transforms import build_transforms
|
|
|
|
|
|
def tile_label(text: str, num_classes: int, background: int) -> int:
|
|
boxes = []
|
|
for line in text.splitlines():
|
|
if not line.strip():
|
|
continue
|
|
parts = line.split()
|
|
if len(parts) != 5:
|
|
raise ValueError(f"Malformed YOLO label: {line}")
|
|
cls, cx, cy, w, h = map(float, parts)
|
|
if (not all(math.isfinite(v) for v in (cls, cx, cy, w, h))
|
|
or not cls.is_integer() or not 0 <= cls < num_classes
|
|
or not 0 <= cx <= 1 or not 0 <= cy <= 1 or not 0 < w <= 1 or not 0 < h <= 1):
|
|
raise ValueError(f"Invalid YOLO label: {line}")
|
|
boxes.append((int(cls), w * h))
|
|
return max(boxes, key=lambda b: b[1])[0] if boxes else background
|
|
|
|
|
|
class YoloTileClassificationDataset(FolderClassificationDataset):
|
|
"""Keep the existing YOLO split. Empty labels are normal; largest box wins.
|
|
|
|
No dataset transform may resample the image. The size is checked on every read.
|
|
"""
|
|
|
|
def __init__(self, yaml_path, split, image_size=(128, 128), transforms=(), classes=None):
|
|
path = Path(yaml_path).resolve()
|
|
cfg = yaml.safe_load(path.read_text())
|
|
names = cfg["names"]
|
|
self.classes = ([names.get(i, names.get(str(i))) for i in range(len(names))]
|
|
if isinstance(names, dict) else list(names))
|
|
if not self.classes or any(not isinstance(n, str) for n in self.classes):
|
|
raise ValueError("Invalid YOLO names mapping")
|
|
if classes is not None and list(classes) != self.classes:
|
|
raise ValueError("Classifier class order must match YOLO names")
|
|
background = self.classes.index("normal")
|
|
self.class_to_idx = dict(zip(self.classes, range(len(self.classes)), strict=True))
|
|
self.root = Path(cfg.get("path", path.parent))
|
|
if not self.root.is_absolute():
|
|
self.root = path.parent / self.root
|
|
folder = self.root / cfg[split]
|
|
if not folder.is_dir():
|
|
raise FileNotFoundError(folder)
|
|
relative = folder.relative_to(self.root)
|
|
if relative.parts[0] != "images":
|
|
raise ValueError("YOLO split must be beneath images/")
|
|
labels = self.root / "labels" / Path(*relative.parts[1:])
|
|
allowed = {"ToTensor", "RandomHorizontalFlip", "RandomVerticalFlip", "PerImageStandardize"}
|
|
if any(t["name"] not in allowed for t in transforms):
|
|
raise ValueError("Fixed tiles forbid resampling transforms")
|
|
self.transforms = build_transforms(transforms)
|
|
self.image_size = tuple(image_size)
|
|
self.samples = []
|
|
for image in sorted(folder.rglob("*.png")):
|
|
label = labels / image.relative_to(folder).with_suffix(".txt")
|
|
if not label.is_file():
|
|
raise FileNotFoundError(f"Missing label (not an implicit background): {label}")
|
|
self.samples.append((image, tile_label(label.read_text(), len(self.classes), background)))
|
|
if not self.samples:
|
|
raise ValueError(f"No PNG tiles in {folder}")
|
|
|
|
def __getitem__(self, index):
|
|
path, label = self.samples[index]
|
|
with Image.open(path) as raw:
|
|
if (raw.height, raw.width) != self.image_size:
|
|
raise ValueError(f"{path}: size {raw.size}, expected {self.image_size}; resize is forbidden")
|
|
image = raw.convert("RGB")
|
|
sample = self.transforms({"image": image})
|
|
return {"image": sample["image"], "label": label,
|
|
"meta": {"path": str(path), "index": index, "class": self.classes[label]}}
|