17 lines
491 B
Python
17 lines
491 B
Python
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()
|