Spaces:
Running on Zero
Running on Zero
File size: 3,362 Bytes
b701455 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | import torch
from unittest.mock import patch
from src.Processors.AutoHDRProcessor import AutoHDRProcessor
def _make_image(batch=1, h=64, w=64, c=3):
# Image expected shape for the processor (B, H, W, C)
return torch.rand(batch, h, w, c)
def test_autohdr_fallback_on_profile_error(monkeypatch):
img = _make_image()
# Simulate ImageCms.profileToProfile raising an error
import src.AutoHDR.ahdr as ahdr_module
def fake_profile_to_profile(*args, **kwargs):
raise OSError("cannot build transform")
monkeypatch.setattr('PIL.ImageCms.profileToProfile', fake_profile_to_profile)
# Should not raise and should return a tensor of same batch shape
out = AutoHDRProcessor.apply(img, ctx=None)
assert isinstance(out, torch.Tensor)
assert out.shape[0] == img.shape[0]
def test_autohdr_detects_lcms_missing_on_import(monkeypatch, caplog):
"""When ImageCms.profileToProfile is unavailable at import time, module sets _HAVE_LCMS=False and logs a warning."""
import importlib
from PIL import ImageCms as ImageCmsPIL
import src.AutoHDR.ahdr as ahdr_module
orig = ImageCmsPIL.profileToProfile
try:
monkeypatch.setattr(ImageCmsPIL, 'profileToProfile', lambda *a, **k: (_ for _ in ()).throw(OSError('cannot build transform')))
caplog.set_level('WARNING')
importlib.reload(ahdr_module)
assert ahdr_module._HAVE_LCMS is False
assert any('LCMS profile transform not available' in rec.getMessage() for rec in caplog.records)
finally:
# Restore clean module state for other tests
monkeypatch.setattr(ImageCmsPIL, 'profileToProfile', orig)
importlib.reload(ahdr_module)
def test_autohdr_handles_rgba_image(monkeypatch):
"""AutoHDR should accept RGBA images and preserve the alpha channel."""
img = _make_image(batch=1, h=64, w=64, c=4)
out = AutoHDRProcessor.apply(img, ctx=None)
assert isinstance(out, torch.Tensor)
assert out.shape[0] == img.shape[0]
# Ensure output still has 4 channels
assert out.shape[-1] == 4
def test_autohdr_rgba_fallback_on_profile_error(monkeypatch):
"""When ImageCms.profileToProfile raises for RGBA input, AutoHDR should fallback without raising."""
img = _make_image(batch=1, h=64, w=64, c=4)
import src.AutoHDR.ahdr as ahdr_module
def fake_profile_to_profile(*args, **kwargs):
raise OSError("cannot build transform")
monkeypatch.setattr('PIL.ImageCms.profileToProfile', fake_profile_to_profile)
out = AutoHDRProcessor.apply(img, ctx=None)
assert isinstance(out, torch.Tensor)
assert out.shape[0] == img.shape[0]
assert out.shape[-1] == 4
def test_autohdr_disables_lcms_on_runtime_failure(monkeypatch):
"""When ImageCms.profileToProfile fails at runtime, AutoHDR should set _HAVE_LCMS=False."""
import src.AutoHDR.ahdr as ahdr_module
# Simulate LCMS being available at import time
ahdr_module._HAVE_LCMS = True
# But make profileToProfile raise during runtime
def fake_profile_to_profile(*args, **kwargs):
raise OSError("cannot build transform")
monkeypatch.setattr('PIL.ImageCms.profileToProfile', fake_profile_to_profile)
img = _make_image()
out = AutoHDRProcessor.apply(img, ctx=None)
assert isinstance(out, torch.Tensor)
assert ahdr_module._HAVE_LCMS is False
|