48 lines
1.5 KiB
Python
48 lines
1.5 KiB
Python
"""
|
|
This unit test requires the dog_and_cat dataset to be set up under
|
|
project_root/datasets/.
|
|
"""
|
|
import os
|
|
import sys
|
|
import unittest
|
|
|
|
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
|
|
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
|
|
sys.path.append(MODULE_ROOT)
|
|
from pytorch_wrapper.ptv_image_folder_dataset import PtvImageFolderDataset
|
|
from pytorch_wrapper.ptv_transforms_to_tensor import PtvTransformsToTensor
|
|
|
|
|
|
class TestPtvImageDataset(unittest.TestCase):
|
|
def setUp(self):
|
|
"""Set up test data."""
|
|
self.node = PtvImageFolderDataset()
|
|
self.compose_node = PtvTransformsToTensor()
|
|
|
|
def test_with_tensor_transform(self):
|
|
"""Tests downloading with transform"""
|
|
transform = self.compose_node.f()[0]
|
|
dc = self.node.f("dog_and_cat/train",
|
|
transform=transform)[0]
|
|
expected_classes = ["cats", "dogs"]
|
|
expected_dict = [("cats", 0), ("dogs", 1)]
|
|
for i, c in enumerate(dc.classes):
|
|
self.assertEqual(c, expected_classes[i])
|
|
|
|
for i, (k, w) in enumerate(dc.class_to_idx.items()):
|
|
self.assertEqual(k, expected_dict[i][0])
|
|
self.assertEqual(w, expected_dict[i][1])
|
|
|
|
# Read data
|
|
data = next(iter(dc))
|
|
|
|
# Check data shape
|
|
self.assertEqual(data[0].size(), (3, 512, 512))
|
|
|
|
# Check class
|
|
self.assertTrue(data[1] == 0) # cat
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|