def test_encode_png(self):
for img_path in get_images(IMAGE_DIR, '.png'):
pil_image = Image.open(img_path)
img_pil = torch.from_numpy(np.array(pil_image))
img_pil = img_pil.permute(2, 0, 1)
png_buf = encode_png(img_pil, compression_level=6)
rec_img = Image.open(io.BytesIO(bytes(png_buf.tolist())))
rec_img = torch.from_numpy(np.array(rec_img))
rec_img = rec_img.permute(2, 0, 1)
self.assertTrue(img_pil.equal(rec_img))
with self.assertRaisesRegex(
RuntimeError, "Input tensor dtype should be uint8"):
encode_png(torch.empty((3, 100, 100), dtype=torch.float32))
with self.assertRaisesRegex(
RuntimeError, "Compression level should be between 0 and 9"):
encode_png(torch.empty((3, 100, 100), dtype=torch.uint8),
compression_level=-1)
with self.assertRaisesRegex(
RuntimeError, "Compression level should be between 0 and 9"):
encode_png(torch.empty((3, 100, 100), dtype=torch.uint8),
compression_level=10)
with self.assertRaisesRegex(
RuntimeError, "The number of channels should be 1 or 3, got: 5"):
encode_png(torch.empty((5, 100, 100), dtype=torch.uint8))
def test_encode_png(img_path):
pil_image = Image.open(img_path)
img_pil = torch.from_numpy(np.array(pil_image))
img_pil = img_pil.permute(2, 0, 1)
png_buf = encode_png(img_pil, compression_level=6)
rec_img = Image.open(io.BytesIO(bytes(png_buf.tolist())))
rec_img = torch.from_numpy(np.array(rec_img))
rec_img = rec_img.permute(2, 0, 1)
assert_equal(img_pil, rec_img)
def test_encode_png_errors():
with pytest.raises(RuntimeError, match="Input tensor dtype should be uint8"):
encode_png(torch.empty((3, 100, 100), dtype=torch.float32))
with pytest.raises(RuntimeError, match="Compression level should be between 0 and 9"):
encode_png(torch.empty((3, 100, 100), dtype=torch.uint8),
compression_level=-1)
with pytest.raises(RuntimeError, match="Compression level should be between 0 and 9"):
encode_png(torch.empty((3, 100, 100), dtype=torch.uint8),
compression_level=10)
with pytest.raises(RuntimeError, match="The number of channels should be 1 or 3, got: 5"):
encode_png(torch.empty((5, 100, 100), dtype=torch.uint8))
