origin

tests/test_handlers.py

@@ -0,0 +1,16 @@
+import unittest
+import torch
+from PIL import Image
+from digit_depth.handlers import image
+
+class Handler(unittest.TestCase):
+    """Test for various data handlers"""
+    def test_tensor_to_PIL(self):
+        instance = image.ImageHandler(Image.open("/home/shuk/digit-depth/images/0001.png"), "RGB")
+        tensor = torch.randn(1, 3, 224, 224)
+        pil_image = instance.tensor_to_PIL()
+        self.assertEqual(pil_image.size, (224, 224))
+
+
+if __name__ == '__main__':
+    unittest.main()