diff --git a/.gitignore b/.gitignore index 3810142f..6183b2ac 100644 --- a/.gitignore +++ b/.gitignore @@ -15,4 +15,5 @@ dist example.py /slurm -wandb \ No newline at end of file +wandb +.DS_Store \ No newline at end of file diff --git a/LION/CTtools/ct_utils.py b/LION/CTtools/ct_utils.py index 3ef4022b..d7372a52 100644 --- a/LION/CTtools/ct_utils.py +++ b/LION/CTtools/ct_utils.py @@ -15,7 +15,7 @@ # AItomotools imports from LION.CTtools.ct_geometry import Geometry -from LION.operators import CTProjectionOp +from LION.operators.CTProjectionOp import CTProjectionOp def from_HU_to_normal(img): diff --git a/LION/classical_algorithms/__init__.py b/LION/classical_algorithms/__init__.py deleted file mode 100644 index 8f478680..00000000 --- a/LION/classical_algorithms/__init__.py +++ /dev/null @@ -1,10 +0,0 @@ -"""LION classical algorithms.""" - -from LION.classical_algorithms.conjugate_gradient import conjugate_gradient -from LION.classical_algorithms.fdk import fdk -from LION.classical_algorithms.fista import fista_l1 -from LION.classical_algorithms.sirt import sirt -from LION.classical_algorithms.spgl1_torch import spgl1_torch -from LION.classical_algorithms.tv_min import tv_min - -__all__ = ["conjugate_gradient", "fdk", "fista_l1", "sirt", "spgl1_torch", "tv_min"] diff --git a/LION/classical_algorithms/fista.py b/LION/classical_algorithms/fista.py index 232b2ac6..81550558 100644 --- a/LION/classical_algorithms/fista.py +++ b/LION/classical_algorithms/fista.py @@ -5,7 +5,7 @@ import torch from tqdm import tqdm -from LION.operators import Operator +from LION.operators.Operator import Operator from LION.utils.math import power_method diff --git a/LION/losses/SUREpgImage.py b/LION/losses/SUREpgImage.py index 995ba215..a814bfdf 100644 --- a/LION/losses/SUREpgImage.py +++ b/LION/losses/SUREpgImage.py @@ -1,42 +1,55 @@ import torch -class SUREpgLoss(): - def __init__(self, zeta: float, sigma2: float, eps1: float = 1e-3, eps2: float = 1e-3, kappa: float = 1.0): +class SUREpgLoss: + def __init__( + self, + zeta: float, + sigma2: float, + eps1: float = 1e-3, + eps2: float = 1e-3, + kappa: float = 1.0, + ): self.zeta = zeta self.sigma2 = sigma2 self.eps1 = eps1 self.eps2 = eps2 self.kappa = kappa - self.p = (1/2) * (1 + self.kappa / (self.kappa**2 + 4)**0.5) + self.p = (1 / 2) * (1 + self.kappa / (self.kappa**2 + 4) ** 0.5) self.q = 1 - self.p - self.a = (self.q / self.p)**0.5 - self.b = (self.p / self.q)**0.5 + self.a = (self.q / self.p) ** 0.5 + self.b = (self.p / self.q) ** 0.5 def __call__(self, model, y): - B=y.shape[0] + B = y.shape[0] N_per_img = y.shape[1] * y.shape[2] * y.shape[3] - fy=model(y) - loss = ((fy - y) ** 2).sum(dim=(1,2,3)) - self.zeta * y.sum(dim=(1,2,3)) - self.sigma2 * N_per_img + fy = model(y) + loss = ( + ((fy - y) ** 2).sum(dim=(1, 2, 3)) + - self.zeta * y.sum(dim=(1, 2, 3)) + - self.sigma2 * N_per_img + ) - #1st derivative MC + # 1st derivative MC delta1 = torch.randn_like(y) fy_perturbated = model(y + self.eps1 * delta1) u = self.zeta * y + self.sigma2 - mc1 = (delta1 * u * (fy_perturbated - fy)).sum(dim=(1,2,3)) + mc1 = (delta1 * u * (fy_perturbated - fy)).sum(dim=(1, 2, 3)) loss += 2.0 * mc1 / self.eps1 - #2nd derivative MC + # 2nd derivative MC u_rand = torch.rand_like(y) - delta2 = torch.where(u_rand < self.p,-self.a * torch.ones_like(y),+self.b * torch.ones_like(y)) + delta2 = torch.where( + u_rand < self.p, -self.a * torch.ones_like(y), +self.b * torch.ones_like(y) + ) - fy_plus = model(y + self.eps2 * delta2) + fy_plus = model(y + self.eps2 * delta2) fy_minus = model(y - self.eps2 * delta2) - mc2 = (delta2 * (fy_plus - 2*fy + fy_minus)).sum(dim=(1,2,3)) + mc2 = (delta2 * (fy_plus - 2 * fy + fy_minus)).sum(dim=(1, 2, 3)) loss -= (2 * self.sigma2 * self.zeta / (self.eps2**2 * self.kappa)) * mc2 - return loss.mean()/N_per_img + return loss.mean() / N_per_img diff --git a/LION/operators/__init__.py b/LION/operators/__init__.py deleted file mode 100644 index 75b3e444..00000000 --- a/LION/operators/__init__.py +++ /dev/null @@ -1,20 +0,0 @@ -"""LION operators.""" - -from LION.operators.CompositeOp import CompositeOp -from LION.operators.CTProjectionOp import CTProjectionOp -from LION.operators.DebiasOp import DebiasOp -from LION.operators.Operator import Operator -from LION.operators.PhotocurrentMapOp import PhotocurrentMapOp, Subsampler -from LION.operators.WalshHadamard2D import WalshHadamard2D -from LION.operators.Wavelet2D import Wavelet2D - -__all__ = [ - "CompositeOp", - "CTProjectionOp", - "DebiasOp", - "Operator", - "PhotocurrentMapOp", - "Subsampler", - "WalshHadamard2D", - "Wavelet2D", -] diff --git a/LION/optimizers/Noisier2Inverse.py b/LION/optimizers/Noisier2Inverse.py index ebaac2b1..2177ac79 100644 --- a/LION/optimizers/Noisier2Inverse.py +++ b/LION/optimizers/Noisier2Inverse.py @@ -16,6 +16,7 @@ import torchvision.transforms.functional as TF + class Noisier2Inverse(LIONsolver): def __init__( self, @@ -39,20 +40,19 @@ def __init__( ) self.operator = ct_utils.make_operator(self.geometry) - self.model.geometry = self.geometry + self.model.geometry = self.geometry self.model.operator = self.operator self.projector = to_autograd(self.operator, num_extra_dims=1) self.recon_fn = self.solver_params.recon_fn - @staticmethod def default_parameters() -> LIONParameter: params = LIONParameter() - params.sigma=3 - params.delta=1 + params.sigma = 3 + params.delta = 1 params.recon_fn = fdk return params - + def mini_batch_step(self, sinos, targets): sigma = self.solver_params.sigma delta = self.solver_params.delta @@ -62,21 +62,20 @@ def mini_batch_step(self, sinos, targets): N = TF.gaussian_blur(N, kernel_size=[ks, ks], sigma=[sigma, sigma]) z = sinos + N - input_recon = self.recon_fn(z,self.model.operator) + input_recon = self.recon_fn(z, self.model.operator) output_recon = self.model(input_recon) output_sino = self.projector(output_recon) target_sino = sinos - N - #Sobolev Loss - #res = output_sino - target_sino - #grad_x = res[:, :, :, 1:] - res[:, :, :, :-1] - #grad_y = res[:, :, 1:, :] - res[:, :, :-1, :] + # Sobolev Loss + # res = output_sino - target_sino + # grad_x = res[:, :, :, 1:] - res[:, :, :, :-1] + # grad_y = res[:, :, 1:, :] - res[:, :, :-1, :] - #batch_loss = ((output_sino - target_sino)**2).mean() + (grad_x**2).mean() + (grad_y**2).mean() - batch_loss= ((output_sino - target_sino)**2).mean() + # batch_loss = ((output_sino - target_sino)**2).mean() + (grad_x**2).mean() + (grad_y**2).mean() + batch_loss = ((output_sino - target_sino) ** 2).mean() return batch_loss - # No validation in Noisier2Inverse def validate(self): return 0 diff --git a/LION/optimizers/Proj2ProjSolver.py b/LION/optimizers/Proj2ProjSolver.py index 02720344..0167ad5b 100644 --- a/LION/optimizers/Proj2ProjSolver.py +++ b/LION/optimizers/Proj2ProjSolver.py @@ -14,6 +14,7 @@ import LION.CTtools.ct_utils as ct_utils from tomosipo.torch_support import to_autograd + class Proj2ProjSolver(LIONsolver): def __init__( self, @@ -37,57 +38,63 @@ def __init__( ) self.operator = ct_utils.make_operator(self.geometry) - self.model.geometry = self.geometry + self.model.geometry = self.geometry self.model.operator = self.operator self.projector = to_autograd(self.operator, num_extra_dims=1) self.recon_fn = self.solver_params.recon_fn - self.global_step=0 + self.global_step = 0 def get_mask(self, shape, step): # shape: (B, C, H, W) mask = torch.ones(shape, device=self.device) - grid = self.solver_params.grid_size - + grid = self.solver_params.grid_size + for b in range(shape[0]): - idx = (step+b) % (grid * grid) + idx = (step + b) % (grid * grid) r = idx // grid c = idx % grid mask[b, :, r::grid, c::grid] = 0 return mask - + def fill_mean(self, sinos, mask): - kernel = torch.tensor([[0, 1, 0], [1, 0, 1], [0, 1, 0]], dtype=torch.float32, device=self.device) / 4.0 - + kernel = ( + torch.tensor( + [[0, 1, 0], [1, 0, 1], [0, 1, 0]], + dtype=torch.float32, + device=self.device, + ) + / 4.0 + ) + kernel = kernel.view(1, 1, 3, 3) local_mean_sino = F.conv2d(sinos, kernel, padding=1) - + filled_sinos = (sinos * mask) + (local_mean_sino * (1 - mask)) return filled_sinos @staticmethod def default_parameters() -> LIONParameter: params = LIONParameter() - params.grid_size = 4 + params.grid_size = 4 params.recon_fn = fdk return params - + def mini_batch_step(self, sinos, targets): - mask=self.get_mask(sinos.shape, self.global_step) + mask = self.get_mask(sinos.shape, self.global_step) self.global_step += sinos.shape[0] - input_sino=self.fill_mean(sinos,mask) - - input_recon=self.recon_fn(input_sino,self.model.operator) + input_sino = self.fill_mean(sinos, mask) + + input_recon = self.recon_fn(input_sino, self.model.operator) output_recon = self.model(input_recon) - output_sino=self.projector(output_recon) + output_sino = self.projector(output_recon) - output_sino_mask=output_sino*(1-mask) - target_sino=sinos*(1-mask) + output_sino_mask = output_sino * (1 - mask) + target_sino = sinos * (1 - mask) batch_loss = ((output_sino_mask - target_sino) ** 2).mean() return batch_loss - # No validation in Proj2Proj def validate(self): return 0 diff --git a/LION/optimizers/Sparse2InverseSolver.py b/LION/optimizers/Sparse2InverseSolver.py index 6099f7f9..3ee083c5 100644 --- a/LION/optimizers/Sparse2InverseSolver.py +++ b/LION/optimizers/Sparse2InverseSolver.py @@ -14,6 +14,7 @@ import LION.CTtools.ct_utils as ct_utils from tomosipo.torch_support import to_autograd + class Sparse2InverseSolver(LIONsolver): def __init__( self, @@ -36,7 +37,7 @@ def __init__( solver_params=solver_params, ) - self.model.geometry = self.geometry + self.model.geometry = self.geometry self.model._make_operator() self.sino_split_count = self.solver_params.sino_split_count self.recon_fn = self.solver_params.recon_fn @@ -44,16 +45,16 @@ def __init__( self._make_sub_operators() @classmethod - def two_two_strategy(cls, sino_split_count) -> list[tuple[int,int]]: - #to return all 2 element combinations from 0 to sino_split_count-1 - combos = [] - for i in range(sino_split_count): - for j in range(i + 1, sino_split_count): - combos.append((i, j)) - return combos - + def two_two_strategy(cls, sino_split_count) -> list[tuple[int, int]]: + # to return all 2 element combinations from 0 to sino_split_count-1 + combos = [] + for i in range(sino_split_count): + for j in range(i + 1, sino_split_count): + combos.append((i, j)) + return combos + def _make_sub_operators(self) -> list[ts.Operator.Operator]: - self.sub_ops = [] + self.sub_ops = [] angles = self.geometry.angles.copy() n = len(angles) k = self.sino_split_count @@ -68,22 +69,24 @@ def _make_sub_operators(self) -> list[ts.Operator.Operator]: for idx_group in range(k): sub_geom = Geometry( - image_shape=tuple(self.geometry.image_shape), # tupla, ints - image_size=tuple(self.geometry.image_size), # tupla, floats - angles=[angles[i] for i in self.subgroup_indices[idx_group]], # list, floats - voxel_size=tuple(self.geometry.voxel_size), # tupla, floats + image_shape=tuple(self.geometry.image_shape), # tupla, ints + image_size=tuple(self.geometry.image_size), # tupla, floats + angles=[ + angles[i] for i in self.subgroup_indices[idx_group] + ], # list, floats + voxel_size=tuple(self.geometry.voxel_size), # tupla, floats mode=self.geometry.mode, dso=float(self.geometry.dso), dsd=float(self.geometry.dsd), - detector_shape=tuple(self.geometry.detector_shape), # tupla, ints - detector_size=tuple(self.geometry.detector_size), # tupla, floats - pixel_size=tuple(self.geometry.pixel_size), # tupla, floats - image_pos=tuple(self.geometry.image_pos) # tupla, floats + detector_shape=tuple(self.geometry.detector_shape), # tupla, ints + detector_size=tuple(self.geometry.detector_size), # tupla, floats + pixel_size=tuple(self.geometry.pixel_size), # tupla, floats + image_pos=tuple(self.geometry.image_pos), # tupla, floats ) sub_op = ct_utils.make_operator(sub_geom) self.sub_ops.append(sub_op) - self.combo_ops_autograd = {} + self.combo_ops_autograd = {} for combo in self.split_combinations: combo_angles = [] for split_idx in combo: @@ -100,7 +103,7 @@ def _make_sub_operators(self) -> list[ts.Operator.Operator]: detector_shape=tuple(self.geometry.detector_shape), detector_size=tuple(self.geometry.detector_size), pixel_size=tuple(self.geometry.pixel_size), - image_pos=tuple(self.geometry.image_pos) + image_pos=tuple(self.geometry.image_pos), ) combo_op = ct_utils.make_operator(combo_geom) self.combo_ops_autograd[combo] = to_autograd(combo_op, num_extra_dims=1) @@ -117,14 +120,13 @@ def _calculate_noisy_sub_recons(self, sinos): subgroup_recons[combo] = mean_subgroup_recon return subgroup_recons - @staticmethod def default_parameters() -> LIONParameter: params = LIONParameter() params.sino_split_count = 4 params.recon_fn = fdk return params - + def mini_batch_step(self, sinos, targets): batch_size = sinos.shape[0] subgroup_recons = self._calculate_noisy_sub_recons(sinos) @@ -132,16 +134,23 @@ def mini_batch_step(self, sinos, targets): total_pixels = 0 for combo, mean_recon in subgroup_recons.items(): output_recon = self.model(mean_recon) - remaining_splits = [i for i in range(self.sino_split_count) if i not in combo] + remaining_splits = [ + i for i in range(self.sino_split_count) if i not in combo + ] projector_combo = tuple(sorted(remaining_splits)) projector = self.combo_ops_autograd[projector_combo] for b in range(batch_size): - projected_sino = projector(output_recon[b:b+1]) - target_sino = torch.cat([sinos[b:b+1, :, self.subgroup_indices[i], :] for i in remaining_splits],dim=2) + projected_sino = projector(output_recon[b : b + 1]) + target_sino = torch.cat( + [ + sinos[b : b + 1, :, self.subgroup_indices[i], :] + for i in remaining_splits + ], + dim=2, + ) batch_loss += self.loss_fn(projected_sino, target_sino) return batch_loss - # No validation in Sparse2Inverse as it is unsupervised learning def validate(self): return 0 @@ -152,6 +161,6 @@ def reconstruct(self, sinos): (sinos.shape[0], *self.geometry.image_shape), device=self.device ) for combo, mean_recon in subgroup_recons.items(): - outputs += self.model(mean_recon) + outputs += self.model(mean_recon) outputs /= len(subgroup_recons) return outputs diff --git a/LION/reconstructors/LIONreconstructor.py b/LION/reconstructors/LIONreconstructor.py index 31a86268..26c5f194 100644 --- a/LION/reconstructors/LIONreconstructor.py +++ b/LION/reconstructors/LIONreconstructor.py @@ -5,17 +5,14 @@ import torch # Import CT utils -from LION.CTtools.ct_geometry import Geometry -from LION.CTtools.ct_utils import make_operator -from LION.models.LIONmodel import to_autograd -from LION.operators import Operator +from LION.operators.Operator import Operator # Base class for a Reconstructor in the LION framework. # This assumes a trained model class LIONReconstructor(ABC): - def __init__(self, operator: Geometry | Operator): + def __init__(self, operator): """ Base class for a Reconstructor in the LION framework. This assumes a trained model. @@ -31,14 +28,21 @@ def __init__(self, operator: Geometry | Operator): if isinstance(operator, Operator): self.geometry = None self.op = operator - elif isinstance(operator, Geometry): - self.geometry = operator - self.op = make_operator(self.geometry) + self.op_autograd = None # TODO: Add support for autograd-wrapped operators that are not CT operators else: - raise ValueError( - "Input operator is neither of class LION.operators.operator.Operator nor LION.CTtools.ct_geometry.Geometry" - ) - self.op_autograd = to_autograd(self.op) + from LION.CTtools.ct_geometry import Geometry + from LION.models.LIONmodel import to_autograd + + if isinstance(operator, Geometry): + from LION.CTtools.ct_utils import make_operator + + self.geometry = operator + self.op = make_operator(self.geometry) + else: + raise ValueError( + "Input operator is neither of class LION.operators.operator.Operator nor LION.CTtools.ct_geometry.Geometry" + ) + self.op_autograd = to_autograd(self.op) def reconstruct(self, sino: torch.Tensor, **kwargs): """ diff --git a/LION/reconstructors/PnP.py b/LION/reconstructors/PnP.py index 9c1d640e..550c7614 100644 --- a/LION/reconstructors/PnP.py +++ b/LION/reconstructors/PnP.py @@ -9,9 +9,6 @@ from tqdm import tqdm from LION.classical_algorithms.conjugate_gradient import conjugate_gradient -from LION.classical_algorithms.fdk import fdk -from LION.CTtools.ct_geometry import Geometry -from LION.operators import Operator from LION.reconstructors.LIONreconstructor import LIONReconstructor from LION.utils.math import power_method @@ -19,7 +16,7 @@ class PnP(LIONReconstructor): def __init__( self, - physics: Geometry | Operator, + physics, prior_fn: Callable[[torch.Tensor], torch.Tensor], algorithm: Literal["ADMM", "HQS", "FBS"] = "ADMM", ): @@ -111,6 +108,8 @@ def hqs_algorithm( :param max_iter: Maximum number of iterations. :return: Reconstructed image tensor. """ + from LION.classical_algorithms.fdk import fdk + if noise_level is not None: print("Warning: ignoring value of mu, estimating from noise_level") sigma = noise_level @@ -192,10 +191,16 @@ def admm_algorithm( v = torch.zeros(self.op.domain_shape, device=measurement.device) u = torch.zeros(self.op.domain_shape, device=measurement.device) + # Normalize data fidelity term by number of measurements so that eta has consistent meaning + # regardless of measurement size + measurement_size = measurement.numel() + def matmul_closure(x: torch.Tensor) -> torch.Tensor: - return self.op.adjoint(self.op(x)) + eta * x + # return self.op.adjoint(self.op(x)) + eta * x + return self.op.adjoint(self.op(x)) / measurement_size + eta * x - AT_y = self.op.adjoint(measurement) + # AT_y = self.op.adjoint(measurement) + AT_y = self.op.adjoint(measurement) / measurement_size iterator = range(max_iter) if prog_bar: iterator = tqdm(iterator, desc="ADMM iterations") diff --git a/LION/reconstructors/__init__.py b/LION/reconstructors/__init__.py deleted file mode 100644 index 22d86201..00000000 --- a/LION/reconstructors/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -"""LION image reconstructors.""" - -from LION.reconstructors.LIONreconstructor import LIONReconstructor -from LION.reconstructors.PnP import PnP - -__all__ = ["LIONReconstructor", "PnP"] diff --git a/env_base.yml b/env_base.yml index a166f4cb..caa82f49 100644 --- a/env_base.yml +++ b/env_base.yml @@ -1,8 +1,33 @@ channels: - - astra-toolbox - conda-forge dependencies: - - astra-toolbox>=2.3 # official `astra-toolbox` channel doesn't ship any `aarch64` build + - astra-toolbox>=2.3 # official `astra-toolbox` channel doesn't ship any `aarch64` build # and `conda-forge` only ships astra-toolbox up to 2.3 - - pip<=25.2 # pip 25.3 drops support for setup.py install which is used to install tomosipo and ts_algorithms from source + - pip<=25.2 # pip 25.3 drops support for setup.py install which is used to install tomosipo and ts_algorithms from source - python=3.12 + - deepinv + - h5py + - imageio + - kornia + - lightning + - matplotlib + - numpy + - opencv + - pandas + - pillow + # - ptwt + - pydicom + # - pylidc + - pytorch + # - pytorch-wavelets + - pywavelets + - scikit-image + - scipy + - spgl1 + # - spyrit + - tifffile + - torchdiffeq + - torchmetrics + - torchvision + - tqdm + - wandb diff --git a/pyproject.toml b/pyproject.toml index 207ae3f0..09b8d565 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,32 +36,32 @@ dependencies = [ # The main dependency is tomosipo which in turns depends on astra-toolbox. # The followings are currently set to prefer newer versions (make sure they don't conflict with the ones above) - "deepinv", # pretrained DRUNet - "h5py", - "imageio", - "kornia", # LION/optimizers/GaussianDenoiserSolver.py - "lightning", - "matplotlib", - "numpy", - "opencv-python", - "pandas", - "pillow", - "pydicom", - "pylidc", + # "deepinv", # pretrained DRUNet + # "h5py", + # "imageio", + # "kornia", # LION/optimizers/GaussianDenoiserSolver.py + # "lightning", + # "matplotlib", + # "numpy", + # "opencv-python", + # "pandas", + # "pillow", "ptwt", # Wavelet transforms for PyTorch, for photocurrent mapping reconstruction baseline - "PyWavelets", # Wavelet transforms for Numpy + # "pydicom", + "pylidc", + # "PyWavelets", # Wavelet transforms for Numpy "pytorch-wavelets", # Wavelet transforms for PyTorch (TODO: switch to ptwt) - "scikit-image", - "scipy", - "spgl1", # SPGL1 solver, for baseline photocurrent mapping reconstruction method + # "scikit-image", + # "scipy", + # "spgl1", # SPGL1 solver, for baseline photocurrent mapping reconstruction method "spyrit", # Fast Walsh-Hadamard Transform - "tifffile", - "torch", - "torchdiffeq", - "torchmetrics", - "torchvision", - "tqdm", - "wandb", + # "tifffile", + # "torch", + # "torchdiffeq", + # "torchmetrics", + # "torchvision", + # "tqdm", + # "wandb", "LION[codestyle]", ] @@ -112,7 +112,9 @@ filterwarnings = [ ] addopts = "-n auto --dist loadfile --maxprocesses=8" markers = [ - "cuda : Tests only to be run when cuda device is available" # Add argument `-m "not cuda"` to the `pytest` command line to skip tests with this marker + "cuda : Tests only to be run when cuda device is available", # Add argument `-m "not cuda"` to the `pytest` command line to skip tests with this marker + "tomosipo : Tests only to be run when tomosipo is available", # Add argument `-m "not tomosipo"` to the `pytest` command line to skip tests with this marker + # We can exclude multiple markers. E.g. use `-m "not (cuda or tomosipo)"` or `-m "not cuda and not tomosipo"` to skip tests with either marker ] # MyPy section diff --git a/tests/.ruff.toml b/pytests/.ruff.toml similarity index 100% rename from tests/.ruff.toml rename to pytests/.ruff.toml diff --git a/tests/classical_algorithms/test_fista.py b/pytests/classical_algorithms/test_fista.py similarity index 92% rename from tests/classical_algorithms/test_fista.py rename to pytests/classical_algorithms/test_fista.py index c8fdfc03..c2d52a50 100644 --- a/tests/classical_algorithms/test_fista.py +++ b/pytests/classical_algorithms/test_fista.py @@ -2,8 +2,10 @@ import pytest import torch -from LION.classical_algorithms import fista_l1 -from LION.operators import CompositeOp, PhotocurrentMapOp, Subsampler, Wavelet2D +from LION.classical_algorithms.fista import fista_l1 +from LION.operators.CompositeOp import CompositeOp +from LION.operators.PhotocurrentMapOp import PhotocurrentMapOp, Subsampler +from LION.operators.Wavelet2D import Wavelet2D def test_fista_l1() -> None: diff --git a/tests/classical_algorithms/test_spgl1.py b/pytests/classical_algorithms/test_spgl1.py similarity index 81% rename from tests/classical_algorithms/test_spgl1.py rename to pytests/classical_algorithms/test_spgl1.py index f58f00f3..7295103d 100644 --- a/tests/classical_algorithms/test_spgl1.py +++ b/pytests/classical_algorithms/test_spgl1.py @@ -2,7 +2,9 @@ import torch from LION.classical_algorithms.spgl1_torch import spgl1_torch -from LION.operators import CompositeOp, PhotocurrentMapOp, Subsampler, Wavelet2D +from LION.operators.CompositeOp import CompositeOp +from LION.operators.PhotocurrentMapOp import PhotocurrentMapOp, Subsampler +from LION.operators.Wavelet2D import Wavelet2D def test_spgl1() -> None: diff --git a/tests/helper.py b/pytests/helper.py similarity index 97% rename from tests/helper.py rename to pytests/helper.py index bce40a67..4ee5c33b 100644 --- a/tests/helper.py +++ b/pytests/helper.py @@ -1,7 +1,7 @@ """Helper/Utilities for test functions.""" import torch -from LION.operators import Operator +from LION.operators.Operator import Operator def dotproduct_adjointness_test( diff --git a/tests/operators/test_composite_op.py b/pytests/operators/test_composite_op.py similarity index 79% rename from tests/operators/test_composite_op.py rename to pytests/operators/test_composite_op.py index d8ecfcf6..909505f6 100644 --- a/tests/operators/test_composite_op.py +++ b/pytests/operators/test_composite_op.py @@ -1,6 +1,9 @@ import torch -from LION.operators import CompositeOp, PhotocurrentMapOp, Subsampler, Wavelet2D -from tests.helper import dotproduct_adjointness_test + +from LION.operators.CompositeOp import CompositeOp +from LION.operators.PhotocurrentMapOp import PhotocurrentMapOp, Subsampler +from LION.operators.Wavelet2D import Wavelet2D +from pytests.helper import dotproduct_adjointness_test def test_composite_op_adjointness(): diff --git a/tests/operators/test_ct_op.py b/pytests/operators/test_ct_op.py similarity index 76% rename from tests/operators/test_ct_op.py rename to pytests/operators/test_ct_op.py index c2791778..586aa487 100644 --- a/tests/operators/test_ct_op.py +++ b/pytests/operators/test_ct_op.py @@ -1,15 +1,18 @@ """Tests for CT operator.""" +import numpy as np import pytest import torch -from LION.CTtools.ct_geometry import Geometry -from LION.CTtools.ct_utils import make_operator -from tests.helper import dotproduct_adjointness_test -from tomosipo.torch_support import to_autograd +from pytests.helper import dotproduct_adjointness_test +@pytest.mark.tomosipo # Add argument `-m "not tomosipo"` to the `pytest` command line to skip this test def test_ct_autograd_op_forward_and_backward(): """Check that to_autograd wraps the CT operator correctly and supports backprop.""" + from LION.CTtools.ct_geometry import Geometry + from LION.CTtools.ct_utils import make_operator + from tomosipo.torch_support import to_autograd + geometry = Geometry.default_parameters() operator = make_operator(geometry=geometry) autograd_operator = to_autograd(operator) @@ -31,9 +34,12 @@ def test_ct_autograd_op_forward_and_backward(): assert torch.isfinite(input_tensor.grad).all() +@pytest.mark.tomosipo # Add argument `-m "not tomosipo"` to the `pytest` command line to skip this test def test_ct_autograd_op_matches_original_operator(): """Check that the autograd wrapper produces the same output as the original operator.""" - import numpy as np + from LION.CTtools.ct_geometry import Geometry + from LION.CTtools.ct_utils import make_operator + from tomosipo.torch_support import to_autograd geometry = Geometry.default_parameters() operator = make_operator(geometry=geometry) @@ -49,8 +55,12 @@ def test_ct_autograd_op_matches_original_operator(): np.testing.assert_allclose(output_autograd, output_np, rtol=1e-5, atol=1e-5) +@pytest.mark.tomosipo # Add argument `-m "not tomosipo"` to the `pytest` command line to skip this test def test_original_op_does_not_propagate_grad(): """Check that the original operator output does not require gradients and backward fails.""" + from LION.CTtools.ct_geometry import Geometry + from LION.CTtools.ct_utils import make_operator + geometry = Geometry.default_parameters() operator = make_operator(geometry=geometry) @@ -70,8 +80,12 @@ def test_original_op_does_not_propagate_grad(): output_tensor.mean().backward() +@pytest.mark.tomosipo # Add argument `-m "not tomosipo"` to the `pytest` command line to skip this test def test_ct_op_adjointness(): """Test CT operator adjoint property.""" + from LION.CTtools.ct_geometry import Geometry + from LION.CTtools.ct_utils import make_operator + geometry = Geometry.default_parameters() operator = make_operator(geometry=geometry) @@ -93,7 +107,11 @@ def test_ct_op_adjointness(): ) +@pytest.mark.tomosipo # Add argument `-m "not tomosipo"` to the `pytest` command line to skip this test def test_ct_op_backward_compatibility_with_tomosipo(): + from LION.CTtools.ct_geometry import Geometry + from LION.CTtools.ct_utils import make_operator + geometry = Geometry.default_parameters() operator = make_operator(geometry=geometry) diff --git a/tests/operators/test_pcm_op.py b/pytests/operators/test_pcm_op.py similarity index 92% rename from tests/operators/test_pcm_op.py rename to pytests/operators/test_pcm_op.py index 73f10589..81853eda 100644 --- a/tests/operators/test_pcm_op.py +++ b/pytests/operators/test_pcm_op.py @@ -1,8 +1,9 @@ """Tests for the photocurrent mapping operator.""" import torch -from LION.operators import PhotocurrentMapOp, Subsampler -from tests.helper import dotproduct_adjointness_test + +from LION.operators.PhotocurrentMapOp import PhotocurrentMapOp, Subsampler +from pytests.helper import dotproduct_adjointness_test def test_pcm_autograd_op_forward_and_backward(): diff --git a/tests/operators/test_wavelet.py b/pytests/operators/test_wavelet.py similarity index 77% rename from tests/operators/test_wavelet.py rename to pytests/operators/test_wavelet.py index 500ec5bb..bf351aac 100644 --- a/tests/operators/test_wavelet.py +++ b/pytests/operators/test_wavelet.py @@ -1,6 +1,7 @@ import torch -from LION.operators import Wavelet2D -from tests.helper import dotproduct_adjointness_test + +from LION.operators.Wavelet2D import Wavelet2D +from pytests.helper import dotproduct_adjointness_test def test_wavelet_db4_adjointness(): diff --git a/tests/operators/test_wht.py b/pytests/operators/test_wht.py similarity index 91% rename from tests/operators/test_wht.py rename to pytests/operators/test_wht.py index e3ae0678..2c321dd4 100644 --- a/tests/operators/test_wht.py +++ b/pytests/operators/test_wht.py @@ -1,7 +1,7 @@ import torch from LION.operators.Operator import Operator +from pytests.helper import dotproduct_adjointness_test from spyrit.core.torch import fwht, ifwht -from tests.helper import dotproduct_adjointness_test def test_wht_adjointness(): diff --git a/tests/utils/test_operator_norm.py b/pytests/utils/test_operator_norm.py similarity index 90% rename from tests/utils/test_operator_norm.py rename to pytests/utils/test_operator_norm.py index 8ba0d1e9..a0903b11 100644 --- a/tests/utils/test_operator_norm.py +++ b/pytests/utils/test_operator_norm.py @@ -1,15 +1,12 @@ """Tests for computing the operator norm of operators.""" +import pytest import torch -from LION.CTtools.ct_geometry import Geometry -from LION.CTtools.ct_utils import make_operator -from LION.operators import ( - Operator, - PhotocurrentMapOp, - Subsampler, - WalshHadamard2D, - Wavelet2D, -) + +from LION.operators.Operator import Operator +from LION.operators.PhotocurrentMapOp import PhotocurrentMapOp, Subsampler +from LION.operators.WalshHadamard2D import WalshHadamard2D +from LION.operators.Wavelet2D import Wavelet2D from LION.utils.math import power_method @@ -55,8 +52,11 @@ def range_shape(self): torch.testing.assert_close(op_norm_computed, op_norm_expected, atol=1e-6, rtol=1e-6) +@pytest.mark.tomosipo # Add argument `-m "not tomosipo"` to the `pytest` command line to skip this test def test_ct_operator_norm_torch(): """Test with CT operator using default geometry.""" + from LION.CTtools.ct_geometry import Geometry + from LION.CTtools.ct_utils import make_operator geometry = Geometry.default_parameters() ct_op = make_operator(geometry) @@ -69,7 +69,6 @@ def test_ct_operator_norm_torch(): def test_pcm_operator_norm_torch(): """Test with photocurrent mapping operator with undersampling.""" - J = 4 N = 1 << J # 16x16 image delta = 1.0 / 4 # subsampling factor, keep only 1/4 of measurements @@ -101,7 +100,6 @@ def test_pcm_operator_norm_torch(): def test_wht_operator_norm_torch(): """Test with Walsh-Hadamard Transform operator.""" - J = 4 N = 1 << J # 16x16 image wht_op = WalshHadamard2D(height=N, width=N) @@ -123,7 +121,6 @@ def test_wht_operator_norm_torch(): def test_wavelet_operator_norm_torch(): """Test with Daubechies 4 wavelet transform operator.""" - image_shape = (16, 16) wavelet_op = Wavelet2D(image_shape, wavelet_name="db4") diff --git a/scripts/example_scripts/PnP_ADMM.py b/scripts/example_scripts/PnP_ADMM.py index 4adedd20..c585e749 100644 --- a/scripts/example_scripts/PnP_ADMM.py +++ b/scripts/example_scripts/PnP_ADMM.py @@ -20,9 +20,9 @@ from tqdm import tqdm # LION imports -from LION.classical_algorithms import fdk +from LION.classical_algorithms.fdk import fdk from LION.experiments import ct_experiments -from LION.reconstructors import PnP +from LION.reconstructors.PnP import PnP # %% [markdown] # ## Setup device diff --git a/scripts/example_scripts/Sparse2Inverse.py b/scripts/example_scripts/Sparse2Inverse.py index b72249bd..8e25d367 100644 --- a/scripts/example_scripts/Sparse2Inverse.py +++ b/scripts/example_scripts/Sparse2Inverse.py @@ -28,7 +28,9 @@ torch.cuda.set_device(device) # Define your data paths -savefolder = pathlib.Path("/store/LION/ea692/LION/LION/trained_models/Sparse2Inverse/Train/SparseAngleLowDoseCTRecon") +savefolder = pathlib.Path( + "/store/LION/ea692/LION/LION/trained_models/Sparse2Inverse/Train/SparseAngleLowDoseCTRecon" +) # Creates the folders if they does not exist savefolder.mkdir(parents=True, exist_ok=True) final_result_fname = "S2I.pt" @@ -37,7 +39,7 @@ # Define experiment experiment = ct_experiments.SparseAngleLowDoseCTRecon() train_dataset = experiment.get_training_dataset() -#30 sinograms for the experiment +# 30 sinograms for the experiment indices = torch.arange(30) train_dataset = data_utils.Subset(train_dataset, indices) @@ -52,7 +54,7 @@ optimizer = Adam(model.parameters(), lr=1e-4) loss_fn = nn.MSELoss() -#Sparse2InverseSolver. +# Sparse2InverseSolver. s2i_params = Sparse2InverseSolver.default_parameters() # Sparse to inverse requires certain user specifications. s2i_params.sino_split_count = 4 @@ -79,7 +81,9 @@ solver.clean_checkpoints() # Test using the training data -savefolder = pathlib.Path("/home/ea692/LION/LION/trained_models/Sparse2Inverse/Test/SparseAngleLowDoseCTRecon/SparseVSNoise/30sin2000ep/64Angles_Haarpsi_and_SSIM") +savefolder = pathlib.Path( + "/home/ea692/LION/LION/trained_models/Sparse2Inverse/Test/SparseAngleLowDoseCTRecon/SparseVSNoise/30sin2000ep/64Angles_Haarpsi_and_SSIM" +) savefolder.mkdir(parents=True, exist_ok=True) model.eval() @@ -87,7 +91,7 @@ solver_params.sino_split_count = 4 solver_params.recon_fn = fdk optimizer = Adam(model.parameters()) -#Not used directly, the solver defines its own loss. +# Not used directly, the solver defines its own loss. loss_fn = nn.MSELoss() solver_sparse = Sparse2InverseSolver( @@ -100,22 +104,24 @@ device=device, ) -#Normalization in order to ensure a fair comparison of structural and perceptual image quality. -def normalize_01(x,y): - x = (x - y.min())/ (y.max() - y.min()) - x[x>1]=1 - x[x<0]=0 +# Normalization in order to ensure a fair comparison of structural and perceptual image quality. +def normalize_01(x, y): + x = (x - y.min()) / (y.max() - y.min()) + x[x > 1] = 1 + x[x < 0] = 0 return x -#SSIM metric + +# SSIM metric def my_ssim(x, y): x = x.detach().squeeze().cpu() y = y.detach().squeeze().cpu() - - target_n = normalize_01(y,y) - sparse_n = normalize_01(x,y) + + target_n = normalize_01(y, y) + sparse_n = normalize_01(x, y) return ssim(target_n, sparse_n, data_range=1) - + + model.eval() solver.set_testing(dataloader, my_ssim) solver.test() diff --git a/scripts/example_scripts/photocurrent_mapping.py b/scripts/example_scripts/photocurrent_mapping.py index 34d1ed0f..3a1c094b 100644 --- a/scripts/example_scripts/photocurrent_mapping.py +++ b/scripts/example_scripts/photocurrent_mapping.py @@ -30,7 +30,13 @@ # %% import torch -device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") +device = torch.device( + "mps" + if torch.backends.mps.is_available() + else "cuda" + if torch.cuda.is_available() + else "cpu" +) torch.set_default_device(device) # %% [markdown] @@ -47,15 +53,15 @@ import matplotlib.pyplot as plt import numpy as np from jaxtyping import Float -from torchmetrics import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure +from torchmetrics.image import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure -from LION.classical_algorithms import fista_l1 +# LION imports +from LION.classical_algorithms.fista import fista_l1 from LION.classical_algorithms.spgl1_torch import spgl1_torch -from LION.operators import CompositeOp, Wavelet2D +from LION.operators.CompositeOp import CompositeOp from LION.operators.DebiasOp import debias_ls - -# LION imports from LION.operators.PhotocurrentMapOp import PhotocurrentMapOp, Subsampler +from LION.operators.Wavelet2D import Wavelet2D from LION.reconstructors.PnP import PnP GrayscaleImage2D = Float[torch.Tensor, "height width"] @@ -69,7 +75,7 @@ # experiments. # %% -data_dir = Path("/home/t/Documents/GIT/LION/data/photocurrent_data") +data_dir = Path("data/photocurrent_data") # data_dir = Path("your/path/to/photocurrent_data") assert data_dir.exists(), f"Data directory {data_dir} does not exist." @@ -134,10 +140,10 @@ def run_pcm_demo( subtract_from_J: int = 1, delta_divided_by: int = 4, log_dir: Path | str = ".", - device: torch.device | str = "cuda:0", + device: torch.device | str | None = None, ): N = 1 << J - im_tensor = torch.tensor(ground_truth_image).unsqueeze(0).unsqueeze(0) # (1,1,H,W) + im_tensor = ground_truth_image.clone().unsqueeze(0).unsqueeze(0) # (1,1,H,W) coarseJ = J - subtract_from_J delta = 1.0 / delta_divided_by @@ -165,8 +171,8 @@ def run_pcm_demo( ) data_range = (im_tensor.max() - im_tensor.min()).item() - psnr = PeakSignalNoiseRatio(data_range=data_range).to(device) - ssim = StructuralSimilarityIndexMeasure(data_range=data_range).to(device) + psnr = PeakSignalNoiseRatio(data_range=data_range).to(device=device) + ssim = StructuralSimilarityIndexMeasure(data_range=data_range).to(device=device) psnr_zero_filled = psnr(zero_filled_recon_tensor, im_tensor) psnr_recon = psnr(recon_tensor, im_tensor) @@ -198,7 +204,9 @@ def run_pcm_demo( # reconstruction methods. # %% -cigs_raw_data: GrayscaleImage2D = np.load(data_dir / cigs_filename) +cigs_raw_data: GrayscaleImage2D = torch.tensor( + np.load(data_dir / cigs_filename), dtype=torch.float32 +) print(f"CIGS data shape: {cigs_raw_data.shape}") # %% [markdown] @@ -256,9 +264,9 @@ def denoiser_fn(x: GrayscaleImage2D) -> GrayscaleImage2D: def run_pnp_admm( pcm_op: PhotocurrentMapOp, pcm_measurement: Measurement1D ) -> GrayscaleImage2D: - admm_iterations = 100 - admm_step_size = 1e5 - cg_max_iter = 100 + admm_iterations = 50 + admm_eta = 0.01 + cg_max_iter = 20 cg_tol = 1e-7 print( @@ -268,7 +276,7 @@ def run_pnp_admm( pnp = PnP(physics=pcm_op, prior_fn=denoiser_fn, algorithm="ADMM") return pnp.admm_algorithm( measurement=pcm_measurement, - eta=admm_step_size, + eta=admm_eta, max_iter=admm_iterations, cg_max_iter=cg_max_iter, cg_tol=cg_tol, @@ -295,8 +303,8 @@ def run_pnp_admm( # # Although PnP-ADMM substantially improves PSNR and SSIM, it can smooth out # fine-scale structures. In the context of defect detection, these small -# features can be crucial, so high PSNR and SSIM alone are not sufficient to -# guarantee that the reconstruction is fit for purpose. +# features can be crucial, so high PSNR and SSIM alone are not always +# sufficient to guarantee that the reconstruction is fit for purpose. # # In the next sections, two compressed sensing baselines with a wavelet # sparsity prior are explored and compared to the PnP-ADMM result. @@ -330,7 +338,7 @@ def run_fista_l1( ) -> GrayscaleImage2D: lam = 10 # Good for Daubechies 4 wavelet transform - max_iter = 1000 + max_iter = 500 tol = 1e-5 debias_max_iter = 10 # TODO: Debiasing seems to not make much difference here @@ -406,7 +414,7 @@ def run_spgl1( ) -> GrayscaleImage2D: lam: float = 1e-3 - max_iter = 1000 + max_iter = 500 tol: float = 1e-4 debias_max_iter = 10 debias_support_tol = 1e-5 diff --git a/tests/__init__.py b/tests/__init__.py deleted file mode 100644 index e69de29b..00000000