41 lines
1.8 KiB
Python
41 lines
1.8 KiB
Python
import pytest
|
|
from PIL import Image
|
|
|
|
from datasets.tile_classification import YoloTileClassificationDataset, tile_label
|
|
from datasets.yolo import write_data_yaml
|
|
|
|
|
|
def test_label_uses_human_class_ids_and_empty_background():
|
|
assert tile_label("", 9, 4) == 4
|
|
assert tile_label("1 0.5 0.5 0.2 0.2", 9, 4) == 1
|
|
assert tile_label("1 0.5 0.5 0.2 0.2\n3 0.5 0.5 0.4 0.4", 9, 4) == 3
|
|
with pytest.raises(ValueError):
|
|
tile_label("1 0.5 0.5 nan 0.2", 9, 4)
|
|
|
|
|
|
def test_existing_split_is_preserved_and_no_resize(tmp_path):
|
|
for split in ("train", "val"):
|
|
(tmp_path / "images" / split).mkdir(parents=True)
|
|
(tmp_path / "labels" / split).mkdir(parents=True)
|
|
Image.new("RGB", (16, 16)).save(tmp_path / "images" / split / "sample.png")
|
|
(tmp_path / "labels" / split / "sample.txt").write_text(
|
|
"1 0.5 0.5 0.2 0.2" if split == "train" else "")
|
|
yaml = write_data_yaml(tmp_path / "data.yaml", ["normal", "blackdot"], root=tmp_path)
|
|
kwargs = {"yaml_path": yaml, "image_size": [16, 16], "transforms": [{"name": "ToTensor"}]}
|
|
train = YoloTileClassificationDataset(split="train", **kwargs)
|
|
val = YoloTileClassificationDataset(split="val", **kwargs)
|
|
assert train.class_counts() == [0, 1]
|
|
assert val.class_counts() == [1, 0]
|
|
assert train[0]["image"].shape == (3, 16, 16)
|
|
Image.new("RGB", (12, 16)).save(train.samples[0][0])
|
|
with pytest.raises(ValueError, match="resize is forbidden"):
|
|
train[0]
|
|
|
|
|
|
def test_missing_labels_are_not_silent_background(tmp_path):
|
|
(tmp_path / "images" / "train").mkdir(parents=True)
|
|
Image.new("RGB", (16, 16)).save(tmp_path / "images" / "train" / "sample.png")
|
|
yaml = write_data_yaml(tmp_path / "data.yaml", ["normal", "blackdot"], root=tmp_path)
|
|
with pytest.raises(FileNotFoundError, match="Missing label"):
|
|
YoloTileClassificationDataset(yaml, "train")
|