super-gradients/tests/unit_tests/cityscapes_dataset_test.py

import unittest

import pkg_resources
import yaml
from torch.utils.data import DataLoader

from super_gradients.training.dataloaders.dataloaders import (
    cityscapes_train,
    cityscapes_val,
    cityscapes_stdc_seg50_train,
    cityscapes_stdc_seg50_val,
    cityscapes_stdc_seg75_val,
    cityscapes_ddrnet_train,
    cityscapes_regseg48_val,
    cityscapes_regseg48_train,
    cityscapes_ddrnet_val,
    cityscapes_stdc_seg75_train,
)
from super_gradients.training.datasets.segmentation_datasets.cityscape_segmentation import CityscapesDataset


class CityscapesDatasetTest(unittest.TestCase):
    def setUp(self) -> None:
        default_config_path = pkg_resources.resource_filename("super_gradients.recipes", "dataset_params/cityscapes_dataset_params.yaml")
        with open(default_config_path, "r") as file:
            self.recipe = yaml.safe_load(file)

    def dataloader_tester(self, dl: DataLoader):
        self.assertTrue(isinstance(dl, DataLoader))
        self.assertTrue(isinstance(dl.dataset, CityscapesDataset))
        it = iter(dl)
        for _ in range(10):
            next(it)

    def test_train_dataset_creation(self):
        train_dataset = CityscapesDataset(**self.recipe["train_dataset_params"])
        for i in range(10):
            image, mask = train_dataset[i]

    def test_val_dataset_creation(self):
        val_dataset = CityscapesDataset(**self.recipe["val_dataset_params"])
        for i in range(10):
            image, mask = val_dataset[i]

    def test_cityscapes_train_dataloader(self):
        dl_train = cityscapes_train()
        self.dataloader_tester(dl_train)

    def test_cityscapes_val_dataloader(self):
        dl_val = cityscapes_val()
        self.dataloader_tester(dl_val)

    def test_cityscapes_stdc_seg50_train_dataloader(self):
        dl_train = cityscapes_stdc_seg50_train()
        self.dataloader_tester(dl_train)

    def test_cityscapes_stdc_seg50_val_dataloader(self):
        dl_val = cityscapes_stdc_seg50_val()
        self.dataloader_tester(dl_val)

    def test_cityscapes_stdc_seg75_train_dataloader(self):
        dl_train = cityscapes_stdc_seg75_train()
        self.dataloader_tester(dl_train)

    def test_cityscapes_stdc_seg75_val_dataloader(self):
        dl_val = cityscapes_stdc_seg75_val()
        self.dataloader_tester(dl_val)

    def test_cityscapes_regseg48_train_dataloader(self):
        dl_train = cityscapes_regseg48_train()
        self.dataloader_tester(dl_train)

    def test_cityscapes_regseg48_val_dataloader(self):
        dl_val = cityscapes_regseg48_val()
        self.dataloader_tester(dl_val)

    def test_cityscapes_ddrnet_train_dataloader(self):
        dl_train = cityscapes_ddrnet_train()
        self.dataloader_tester(dl_train)

    def test_cityscapes_ddrnet_val_dataloader(self):
        dl_val = cityscapes_ddrnet_val()
        self.dataloader_tester(dl_val)


if __name__ == "__main__":
    unittest.main()