File size: 897 Bytes
6a39ae1 a56a51f 6a39ae1 |
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 |
import os
import shutil
import sys
from pathlib import Path
project_root = Path(__file__).resolve().parent.parent.parent
sys.path.append(str(project_root))
from yolo.config.config import Config
from yolo.tools.dataset_preparation import prepare_dataset, prepare_weight
def test_prepare_dataset(train_cfg: Config):
dataset_path = Path("tests/data")
if dataset_path.exists():
shutil.rmtree(dataset_path)
prepare_dataset(train_cfg.dataset, task="train")
prepare_dataset(train_cfg.dataset, task="val")
images_path = Path("tests/data/images")
for data_type in images_path.iterdir():
assert len(os.listdir(data_type)) == 5
annotations_path = Path("tests/data/annotations")
assert "instances_val.json" in os.listdir(annotations_path)
assert "instances_train.json" in os.listdir(annotations_path)
def test_prepare_weight():
prepare_weight()
|