40 lines
1.3 KiB
Python
40 lines
1.3 KiB
Python
import unittest
|
|
|
|
import torch
|
|
from PIL import Image
|
|
|
|
from utils.image import ensure_contiguous, pil_to_tensor, tensor_to_pil
|
|
|
|
|
|
class ImageUtilsTests(unittest.TestCase):
|
|
def test_tensor_to_pil_converts_batch_item(self):
|
|
tensor = torch.tensor(
|
|
[[[[0.0, 0.5, 1.0], [1.0, 0.0, 0.0]]]],
|
|
dtype=torch.float32,
|
|
)
|
|
image = tensor_to_pil(tensor, batch_index=0)
|
|
|
|
self.assertEqual(image.size, (2, 1))
|
|
self.assertEqual(image.getpixel((0, 0)), (0, 127, 255))
|
|
self.assertEqual(image.getpixel((1, 0)), (255, 0, 0))
|
|
|
|
def test_pil_to_tensor_adds_channel_for_grayscale(self):
|
|
grayscale = Image.new("L", (3, 2), color=128)
|
|
tensor = pil_to_tensor(grayscale)
|
|
|
|
self.assertEqual(tuple(tensor.shape), (1, 2, 3, 1))
|
|
self.assertAlmostEqual(float(tensor[0, 0, 0, 0]), 128 / 255.0, places=6)
|
|
|
|
def test_ensure_contiguous_makes_non_contiguous_tensor_contiguous(self):
|
|
base = torch.arange(24, dtype=torch.float32).reshape(2, 3, 4)
|
|
non_contiguous = base.transpose(1, 2)
|
|
self.assertFalse(non_contiguous.is_contiguous())
|
|
|
|
contiguous = ensure_contiguous(non_contiguous)
|
|
self.assertTrue(contiguous.is_contiguous())
|
|
self.assertTrue(torch.equal(contiguous, non_contiguous))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|