ai-training/tests/test_dataset_formation.py

import shutil
import sys
import types
from os import path as osp
from pathlib import Path

import pytest

import constants as c_mod


def _stub_train_dependencies():
    if getattr(_stub_train_dependencies, "_done", False):
        return

    def add_mod(name):
        if name in sys.modules:
            return sys.modules[name]
        m = types.ModuleType(name)
        sys.modules[name] = m
        return m

    ultra = add_mod("ultralytics")

    class YOLO:
        pass

    ultra.YOLO = YOLO

    def fake_client(*_a, **_k):
        return types.SimpleNamespace(
            upload_fileobj=lambda *_a, **_k: None,
            download_file=lambda *_a, **_k: None,
        )

    boto = add_mod("boto3")
    boto.client = fake_client
    add_mod("netron")
    add_mod("requests")
    _stub_train_dependencies._done = True


_stub_train_dependencies()


def _prepare_form_dataset(
    monkeypatch,
    tmp_path,
    constants_patch,
    fixture_images_dir,
    fixture_labels_dir,
    count,
    corrupt_stems,
):
    constants_patch(tmp_path)
    import train

    proc_img = Path(c_mod.processed_images_dir)
    proc_lbl = Path(c_mod.processed_labels_dir)
    proc_img.mkdir(parents=True, exist_ok=True)
    proc_lbl.mkdir(parents=True, exist_ok=True)

    imgs = sorted(fixture_images_dir.glob("*.jpg"))[:count]
    for p in imgs:
        stem = p.stem
        shutil.copy2(fixture_images_dir / f"{stem}.jpg", proc_img / f"{stem}.jpg")
        dst = proc_lbl / f"{stem}.txt"
        shutil.copy2(fixture_labels_dir / f"{stem}.txt", dst)
        if stem in corrupt_stems:
            dst.write_text("0 1.5 0.5 0.1 0.1\n", encoding="utf-8")

    today_ds = osp.join(c_mod.datasets_dir, train.today_folder)
    monkeypatch.setattr(train, "today_dataset", today_ds)
    monkeypatch.setattr(train, "processed_images_dir", c_mod.processed_images_dir)
    monkeypatch.setattr(train, "processed_labels_dir", c_mod.processed_labels_dir)
    monkeypatch.setattr(train, "corrupted_images_dir", c_mod.corrupted_images_dir)
    monkeypatch.setattr(train, "corrupted_labels_dir", c_mod.corrupted_labels_dir)
    monkeypatch.setattr(train, "datasets_dir", c_mod.datasets_dir)
    return train


def _count_jpg(p):
    return len(list(Path(p).glob("*.jpg")))


def test_bt_dsf_01_split_ratio_70_20_10(
    monkeypatch,
    tmp_path,
    constants_patch,
    fixture_images_dir,
    fixture_labels_dir,
):
    train = _prepare_form_dataset(
        monkeypatch,
        tmp_path,
        constants_patch,
        fixture_images_dir,
        fixture_labels_dir,
        100,
        set(),
    )
    train.form_dataset()
    base = train.today_dataset
    assert _count_jpg(Path(base, "train", "images")) == 70
    assert _count_jpg(Path(base, "valid", "images")) == 20
    assert _count_jpg(Path(base, "test", "images")) == 10


def test_bt_dsf_02_six_subdirectories(
    monkeypatch,
    tmp_path,
    constants_patch,
    fixture_images_dir,
    fixture_labels_dir,
):
    train = _prepare_form_dataset(
        monkeypatch,
        tmp_path,
        constants_patch,
        fixture_images_dir,
        fixture_labels_dir,
        100,
        set(),
    )
    train.form_dataset()
    base = Path(train.today_dataset)
    assert (base / "train" / "images").is_dir()
    assert (base / "train" / "labels").is_dir()
    assert (base / "valid" / "images").is_dir()
    assert (base / "valid" / "labels").is_dir()
    assert (base / "test" / "images").is_dir()
    assert (base / "test" / "labels").is_dir()


def test_bt_dsf_03_total_files_one_hundred(
    monkeypatch,
    tmp_path,
    constants_patch,
    fixture_images_dir,
    fixture_labels_dir,
):
    train = _prepare_form_dataset(
        monkeypatch,
        tmp_path,
        constants_patch,
        fixture_images_dir,
        fixture_labels_dir,
        100,
        set(),
    )
    train.form_dataset()
    base = train.today_dataset
    n = (
        _count_jpg(Path(base, "train", "images"))
        + _count_jpg(Path(base, "valid", "images"))
        + _count_jpg(Path(base, "test", "images"))
    )
    assert n == 100


def test_bt_dsf_04_corrupted_labels_quarantined(
    monkeypatch,
    tmp_path,
    constants_patch,
    fixture_images_dir,
    fixture_labels_dir,
):
    stems = [p.stem for p in sorted(fixture_images_dir.glob("*.jpg"))[:100]]
    corrupt = set(stems[:5])
    train = _prepare_form_dataset(
        monkeypatch,
        tmp_path,
        constants_patch,
        fixture_images_dir,
        fixture_labels_dir,
        100,
        corrupt,
    )
    train.form_dataset()
    base = train.today_dataset
    split_total = (
        _count_jpg(Path(base, "train", "images"))
        + _count_jpg(Path(base, "valid", "images"))
        + _count_jpg(Path(base, "test", "images"))
    )
    assert split_total == 95
    assert _count_jpg(c_mod.corrupted_images_dir) == 5
    assert len(list(Path(c_mod.corrupted_labels_dir).glob("*.txt"))) == 5


@pytest.mark.resilience
def test_rt_dsf_01_empty_processed_no_crash(
    monkeypatch,
    tmp_path,
    constants_patch,
    fixture_images_dir,
    fixture_labels_dir,
):
    train = _prepare_form_dataset(
        monkeypatch,
        tmp_path,
        constants_patch,
        fixture_images_dir,
        fixture_labels_dir,
        0,
        set(),
    )
    train.form_dataset()
    base = Path(train.today_dataset)
    assert base.is_dir()


@pytest.mark.resource_limit
def test_rl_dsf_01_split_ratios_sum_hundred():
    import train

    assert train.train_set + train.valid_set + train.test_set == 100


@pytest.mark.resource_limit
def test_rl_dsf_02_no_filename_duplication_across_splits(
    monkeypatch,
    tmp_path,
    constants_patch,
    fixture_images_dir,
    fixture_labels_dir,
):
    train = _prepare_form_dataset(
        monkeypatch,
        tmp_path,
        constants_patch,
        fixture_images_dir,
        fixture_labels_dir,
        100,
        set(),
    )
    train.form_dataset()
    base = Path(train.today_dataset)
    names = []
    for split in ("train", "valid", "test"):
        for f in (base / split / "images").glob("*.jpg"):
            names.append(f.name)
    assert len(names) == len(set(names))
    assert len(names) == 100