transform-image-operator-template/test_transform_image_operator_template.py

import unittest
from PIL import Image
from torchvision import transforms
from transform_image_operator_template import TransformImageOperatorTemplate
from config import TEST_IMG, SIZE


class TestTransformImageOperatorTemplate(unittest.TestCase):
    tfms = transforms.Compose(
        [
            transforms.Resize(SIZE),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ]
    )
    img1 = tfms(test_img).unsqueeze(0)

    def test_transform_image(self):
        op = TransformImageOperatorTemplate(SIZE)
        outputs = op(TEST_IMG)
        print("The output tyep of operator:", type(outputs.img_transformed))
        c = (self.img1.numpy() == outputs.img_transformed.numpy())
        self.assertEqual(c.all(), True)


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