117 lines
3.5 KiB
Python
117 lines
3.5 KiB
Python
import numpy as np
|
|
from PIL import Image
|
|
|
|
from smart_metadata_reader import image_io
|
|
|
|
|
|
class FakeTensor:
|
|
def __init__(self, array):
|
|
self.array = np.array(array)
|
|
|
|
@property
|
|
def shape(self):
|
|
return self.array.shape
|
|
|
|
def reshape(self, *shape):
|
|
return FakeTensor(self.array.reshape(*shape))
|
|
|
|
def float(self):
|
|
return FakeTensor(self.array.astype(np.float32))
|
|
|
|
def __truediv__(self, value):
|
|
return FakeTensor(self.array / value)
|
|
|
|
def tolist(self):
|
|
return self.array.tolist()
|
|
|
|
|
|
class FakeTorch:
|
|
float32 = np.float32
|
|
|
|
@staticmethod
|
|
def tensor(data, dtype=None):
|
|
return FakeTensor(np.array(data, dtype=dtype))
|
|
|
|
@staticmethod
|
|
def from_numpy(array):
|
|
return FakeTensor(array)
|
|
|
|
|
|
def use_fake_torch(monkeypatch):
|
|
monkeypatch.setattr(image_io, "_torch_module", lambda: FakeTorch)
|
|
|
|
|
|
def test_rgb_png_outputs_image_and_empty_mask(tmp_path, monkeypatch):
|
|
use_fake_torch(monkeypatch)
|
|
path = tmp_path / "rgb.png"
|
|
Image.new("RGB", (2, 1), (128, 64, 32)).save(path)
|
|
|
|
image, mask, width, height = image_io.load_image_and_mask(path)
|
|
|
|
assert (width, height) == (2, 1)
|
|
assert image.shape == (1, 1, 2, 3)
|
|
assert mask.shape == (1, 1, 2)
|
|
assert np.allclose(mask.array, 0.0)
|
|
assert np.allclose(image.array[0, 0, 0], [128 / 255.0, 64 / 255.0, 32 / 255.0])
|
|
|
|
|
|
def test_rgba_png_outputs_alpha_mask(tmp_path, monkeypatch):
|
|
use_fake_torch(monkeypatch)
|
|
path = tmp_path / "rgba.png"
|
|
image = Image.new("RGBA", (3, 1))
|
|
image.putdata([(255, 0, 0, 255), (0, 255, 0, 0), (0, 0, 255, 128)])
|
|
image.save(path)
|
|
|
|
loaded, mask, width, height = image_io.load_image_and_mask(path)
|
|
|
|
assert (width, height) == (3, 1)
|
|
assert loaded.shape == (1, 1, 3, 3)
|
|
assert mask.shape == (1, 1, 3)
|
|
assert np.allclose(mask.array[0, 0], [0.0, 1.0, 1.0 - (128 / 255.0)])
|
|
|
|
|
|
def test_large_image_conversion_keeps_expected_shape(tmp_path, monkeypatch):
|
|
use_fake_torch(monkeypatch)
|
|
path = tmp_path / "large.png"
|
|
Image.new("RGB", (640, 512), (10, 20, 30)).save(path)
|
|
|
|
image, mask, width, height = image_io.load_image_and_mask(path)
|
|
|
|
assert (width, height) == (640, 512)
|
|
assert image.shape == (1, 512, 640, 3)
|
|
assert mask.shape == (1, 512, 640)
|
|
|
|
|
|
def test_numpy_path_does_not_use_getdata(tmp_path, monkeypatch):
|
|
use_fake_torch(monkeypatch)
|
|
path = tmp_path / "numpy.png"
|
|
Image.new("RGBA", (2, 2), (10, 20, 30, 128)).save(path)
|
|
|
|
def broken_getdata(self):
|
|
raise AssertionError("numpy path should not call getdata")
|
|
|
|
monkeypatch.setattr(Image.Image, "getdata", broken_getdata)
|
|
|
|
image, mask, width, height = image_io.load_image_and_mask(path)
|
|
|
|
assert (width, height) == (2, 2)
|
|
assert image.shape == (1, 2, 2, 3)
|
|
assert mask.shape == (1, 2, 2)
|
|
|
|
|
|
def test_numpy_unavailable_fallback_shape_matches_numpy_path(tmp_path, monkeypatch):
|
|
path = tmp_path / "fallback.png"
|
|
Image.new("RGBA", (4, 3), (10, 20, 30, 128)).save(path)
|
|
|
|
use_fake_torch(monkeypatch)
|
|
numpy_image, numpy_mask, width, height = image_io.load_image_and_mask(path)
|
|
|
|
monkeypatch.setattr(image_io, "_numpy_module", lambda: None, raising=False)
|
|
fallback_image, fallback_mask, fallback_width, fallback_height = image_io.load_image_and_mask(path)
|
|
|
|
assert (width, height) == (fallback_width, fallback_height)
|
|
assert fallback_image.shape == numpy_image.shape
|
|
assert fallback_mask.shape == numpy_mask.shape
|
|
assert np.allclose(fallback_image.array, numpy_image.array)
|
|
assert np.allclose(fallback_mask.array, numpy_mask.array)
|