From 911c76ae02c9d98999188fae82fa6f33b803f47c Mon Sep 17 00:00:00 2001 From: cslht Date: Tue, 4 Aug 2026 17:28:52 -0700 Subject: [PATCH 1/2] Fix PSNR/SSIM crashes on common inputs - psnr/ssim: unbatched mode passed torch tensors to skimage (TypeError); now uses the numpy-converted x_/target_ - ssim: default channel_axis=1 was wrong for LION's channels-first layout; crashed on multichannel input and silently returned wrong values on 2D. Now None, inferred per sample (2D -> None, 3D+ -> 0) - ssim: vals[i] = skim_ssim(...) assigned numpy.float32 into a torch tensor (TypeError); wrapped in torch.tensor(...) - psnr/ssim: reduce=str | None was a default value, not a type hint, so calling without reduce always raised ValueError; now reduce: str | None = None --- LION/metrics/psnr.py | 8 ++++++-- LION/metrics/ssim.py | 26 +++++++++++++++++--------- 2 files changed, 23 insertions(+), 11 deletions(-) diff --git a/LION/metrics/psnr.py b/LION/metrics/psnr.py index 626a399e..c5e4c4f6 100644 --- a/LION/metrics/psnr.py +++ b/LION/metrics/psnr.py @@ -24,7 +24,11 @@ def __init__(self): """ def forward( - self, x: torch.Tensor, target: torch.Tensor, reduce=str | None, batched=True + self, + x: torch.Tensor, + target: torch.Tensor, + reduce: str | None = None, + batched=True, ) -> torch.Tensor: if x.shape != target.shape: raise ShapeMismatchException( @@ -49,5 +53,5 @@ def forward( ) else: return torch.tensor( - skim_psnr(x, target, data_range=target_.max() - target_.min()) + skim_psnr(x_, target_, data_range=target_.max() - target_.min()) ) diff --git a/LION/metrics/ssim.py b/LION/metrics/ssim.py index b86eeac9..a0ba8534 100644 --- a/LION/metrics/ssim.py +++ b/LION/metrics/ssim.py @@ -28,9 +28,9 @@ def forward( self, x: torch.Tensor, target: torch.Tensor, - reduce=str | None, + reduce: str | None = None, batched=True, - channel_axis: int | None = 1, + channel_axis: int | None = None, ) -> torch.Tensor: if x.shape != target.shape: raise ShapeMismatchException( @@ -38,16 +38,22 @@ def forward( ) x_ = x.detach().cpu().numpy().squeeze() target_ = target.detach().cpu().numpy().squeeze() + # LION is channels-first, so a 3D sample means (C, W, H) -> channel axis 0 if batched: # shape either B, C, W, H, ... or B, W, H, ... # if it's not, then that's your fault not mine, you told me it was batched vals = torch.empty((x.shape[0])) for i in range(x.shape[0]): - vals[i] = skim_ssim( - x_[i], - target_[i], - data_range=target_[i].max() - target_[i].min(), - channel_axis=channel_axis, + sample_axis = channel_axis + if sample_axis is None: + sample_axis = 0 if x_[i].ndim >= 3 else None + vals[i] = torch.tensor( + skim_ssim( + x_[i], + target_[i], + data_range=target_[i].max() - target_[i].min(), + channel_axis=sample_axis, + ) ) if reduce is None: @@ -59,10 +65,12 @@ def forward( f"expected one of 'mean' or None for parameter 'reduce', got {reduce}" ) else: + if channel_axis is None: + channel_axis = 0 if x_.ndim >= 3 else None return torch.tensor( skim_ssim( - x, - target, + x_, + target_, data_range=target_.max() - target_.min(), channel_axis=channel_axis, ) From 7c1340e7db13b9ddcac062d416493a02c7e2c302 Mon Sep 17 00:00:00 2001 From: cslht Date: Tue, 4 Aug 2026 19:28:58 -0700 Subject: [PATCH 2/2] Apply black formatting to three files black --all-files in pre-commit flags LIONmodel.py and two classical_algorithms tests (from #194). Formatting only, no behavioural changes. --- LION/models/LIONmodel.py | 6 ++++-- tests/classical_algorithms/test_sirt.py | 7 ++++--- tests/classical_algorithms/test_tv_min.py | 4 +--- 3 files changed, 9 insertions(+), 8 deletions(-) diff --git a/LION/models/LIONmodel.py b/LION/models/LIONmodel.py index 1999baef..6f0f44cc 100644 --- a/LION/models/LIONmodel.py +++ b/LION/models/LIONmodel.py @@ -83,8 +83,10 @@ def __init__( f"Expected geometry to be of type Geometry or None, but got {type(geometry).__name__}. " "If you passed positional arguments to the model, please verify their order matches the model's __init__ signature." ) - - if model_parameters is not None and not isinstance(model_parameters, LIONModelParameter): + + if model_parameters is not None and not isinstance( + model_parameters, LIONModelParameter + ): raise TypeError( f"Expected model_parameters to be of type LIONModelParameter or None, but got {type(model_parameters).__name__}. " "Ensure you are not accidentally passing a Geometry object as model_parameters." diff --git a/tests/classical_algorithms/test_sirt.py b/tests/classical_algorithms/test_sirt.py index ebdf591a..c71280f7 100644 --- a/tests/classical_algorithms/test_sirt.py +++ b/tests/classical_algorithms/test_sirt.py @@ -19,6 +19,7 @@ class Geometry: """Minimal Geometry stub for isinstance checks.""" + pass @@ -35,6 +36,7 @@ class Geometry: class NoDataException(Exception): """Stub matching LION.exceptions.exceptions.NoDataException.""" + pass @@ -62,6 +64,7 @@ class NoDataException(Exception): class MockOp: """Minimal operator stub matching tomosipo's interface.""" + def __init__(self, domain_shape=(1, 16, 16), range_shape=(1, 20, 20)): self.domain_shape = domain_shape self.range_shape = range_shape @@ -71,9 +74,7 @@ def __init__(self, domain_shape=(1, 16, 16), range_shape=(1, 20, 20)): def mock_ts_sirt(): """Replace the real ts_algorithms.sirt with a deterministic mock.""" _sirt_mod.ts_sirt = MagicMock() - _sirt_mod.ts_sirt.side_effect = lambda op, y, *a, **kw: torch.zeros( - op.domain_shape - ) + _sirt_mod.ts_sirt.side_effect = lambda op, y, *a, **kw: torch.zeros(op.domain_shape) yield _sirt_mod.ts_sirt = ts_sirt # restore diff --git a/tests/classical_algorithms/test_tv_min.py b/tests/classical_algorithms/test_tv_min.py index a80d96ab..113f87ad 100644 --- a/tests/classical_algorithms/test_tv_min.py +++ b/tests/classical_algorithms/test_tv_min.py @@ -66,9 +66,7 @@ def __init__(self, domain_shape=(1, 16, 16), range_shape=(1, 20, 20)): def mock_ts_tv_min(): """Replace the real ts_algorithms.tv_min2d with a deterministic mock.""" _tv_mod.ts_tv_min = MagicMock() - _tv_mod.ts_tv_min.side_effect = lambda op, y, *a, **kw: torch.zeros( - op.domain_shape - ) + _tv_mod.ts_tv_min.side_effect = lambda op, y, *a, **kw: torch.zeros(op.domain_shape) yield _tv_mod.ts_tv_min = ts_tv_min