Fix PSNR/SSIM crashes on common inputs

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())
             )

LION/metrics/ssim.py

@@ -28,26 +28,32 @@ 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(
                 f"x (shape {x.shape}) and target (shape {target.shape}) tensors must match to compare with SSIM"
             )
         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,
                 )

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."

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
 

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