# =============================================================================
# This file is part of LION library
# License : GPL-3
#
# Author : Ander Biguri, Alvaro Exposito Martinez
# =============================================================================
from typing import Optional
from LION.models.LIONmodel import LIONmodel
import torch
from tqdm import tqdm
import torch.nn.functional as F
from torch.optim.optimizer import Optimizer
from LION.optimizers.LIONsolver import LIONsolver, warn_if_not_mse
from LION.utils.parameter import LIONParameter, OperatorParameters
import random
from LION.utils.image_transforms import rotate
from LION.utils.citation import citation
def _robust_std(values: torch.Tensor) -> torch.Tensor:
"""Standard deviation of Gaussian values from their median absolute
deviation, insensitive to outliers (here: edges of the object)."""
deviation = (values - values.median()).abs().median()
return 1.4826 * deviation
[docs]
def estimate_noise_std(sinos: torch.Tensor, noise_kernel=None) -> float:
"""Estimate the standard deviation of additive Gaussian noise in sinograms.
The noise is ``noise_kernel * w`` along the detector, with ``w`` iid
Gaussian of the returned standard deviation (``noise_kernel=None``: iid
noise). The estimate uses second differences along the detector,
``y[j-1] - 2 y[j] + y[j+1]``, where a smooth signal cancels: for that noise
their standard deviation is ``std * ||second difference of noise_kernel||``.
A robust standard deviation (from the median absolute deviation) ignores
the large differences at edges of the object. Triplets with a 0, e.g.
post-log values clipped at 0 in air, are left out.
The noise of CT data depends on the signal (it is larger behind more
attenuation), so this is a typical level over the sinograms. The E2I
paper measures it in a background region instead; give the solver's
``noise_std`` to use such a value.
``sinos``: sinograms (..., angles, detectors); ``noise_kernel``: the 1D
blur of the noise along the detector, or None. Returns the standard
deviation before the blur.
"""
left, centre, right = sinos[..., :-2], sinos[..., 1:-1], sinos[..., 2:]
keep = (left != 0) & (centre != 0) & (right != 0)
if keep.sum() < 100:
raise ValueError(
"Too few nonzero detector pixels to estimate the noise; "
"set the solver's noise_std."
)
differences = (left - 2 * centre + right)[keep].flatten().float()
if differences.numel() > 10_000_000: # the median needs no more than this
step = differences.numel() // 10_000_000 + 1
differences = differences[::step]
kernel = torch.ones(1) if noise_kernel is None else torch.as_tensor(noise_kernel)
kernel = kernel.flatten().double()
# second differences of the kernel, with zeros outside it
padded = torch.cat([kernel.new_zeros(2), kernel, kernel.new_zeros(2)])
gain = torch.linalg.vector_norm(padded[:-2] - 2 * padded[1:-1] + padded[2:])
return float(_robust_std(differences) / gain)
[docs]
class Equivariance2InverseSolver(LIONsolver):
"""Equivariance2Inverse (E2I).
The training loss is the mean squared error, as in the method's
formulation: a projection-wise cross-validation (data consistency) term
plus ``equivariance_strength`` (lambda in the paper) times the equivariance
term. ``loss_fn`` is only used as the default validation and testing
function.
The equivariance term rotates the reconstruction by a random angle,
projects it and adds new Gaussian noise to the (post-log) sinogram, as in
the paper and the authors' code (github.com/D1rk123/equivariance2inverse):
iid with standard deviation ``noise_std`` or, with ``noise_kernel``, blurred
along the detector with that kernel (the detector's point response).
``noise_std=None`` estimates it from the first training batch
(``estimate_noise_std``) and stores it in the solver parameters.
"""
def __init__(
self,
model: LIONmodel,
optimizer: Optimizer,
loss_fn,
operator_params: OperatorParameters = None,
solver_params: Optional[LIONParameter] = None,
model_regularization=None,
scheduler=None,
scheduler_step: str = "epoch",
verbose: int = 1,
device: torch.device = None,
seed: Optional[int] = 0,
deterministic: bool = False,
) -> None:
super().__init__(
model,
optimizer,
loss_fn,
operator_params=operator_params,
solver_params=solver_params,
model_regularization=model_regularization,
scheduler=scheduler,
scheduler_step=scheduler_step,
verbose=verbose,
device=device,
seed=seed,
deterministic=deterministic,
)
warn_if_not_mse(loss_fn, "Equivariance2InverseSolver")
self.op_params = self.operator_params
self.projector, _ = self.op.autograd()
[docs]
@staticmethod
def default_parameters() -> LIONParameter:
"""``equivariance_strength`` (lambda), ``noise_std``,
``noise_kernel`` and ``recon_fn``, as below."""
params = LIONParameter()
# lambda in the paper, chosen per dataset there
params.equivariance_strength = 1.0
# Standard deviation of the Gaussian noise added to the rotated
# (post-log) sinograms, before noise_kernel. None: estimated from the
# first training batch (estimate_noise_std).
params.noise_std = None
# 1D kernel (along the detector) that blurs that noise, e.g. a measured
# point response function. None: iid noise (the paper's default).
params.noise_kernel = None
# How measurements are reconstructed, recon_fn(y, op). "default": the
# operator's default reconstruction (FDK for CT), replaced by that
# function when the solver is made.
params.recon_fn = "default"
return params
def _noise(self, y: torch.Tensor) -> torch.Tensor:
"""New Gaussian noise for the sinograms (N_BG in the paper)."""
noise = torch.randn_like(y) * self.solver_params.noise_std
kernel = self.solver_params.noise_kernel
if kernel is not None:
kernel = torch.as_tensor(kernel, dtype=y.dtype, device=y.device)
kernel = kernel.flatten()
# as the authors' Conv2d(padding="same", padding_mode="reflect")
left = (kernel.numel() - 1) // 2
lines = noise.reshape(-1, 1, noise.shape[-1])
lines = F.pad(lines, (left, kernel.numel() - 1 - left), mode="reflect")
noise = F.conv1d(lines, kernel.view(1, 1, -1)).reshape(noise.shape)
return noise
[docs]
def mini_batch_step(self, y, targets):
"""The loss: projection-wise cross-validation (one angle left out,
the model's output projected onto it) plus ``equivariance_strength``
times the equivariance term (``targets`` aren't used:
self-supervised)."""
if self.solver_params.noise_std is None:
self.solver_params.noise_std = estimate_noise_std(
y, self.solver_params.noise_kernel
)
if self.verbose >= 2:
tqdm.write(
"Equivariance2InverseSolver: noise_std estimated from the "
f"data: {self.solver_params.noise_std:.4g}"
)
# masking
NP = y.shape[2]
YJ_num = torch.randint(0, NP, (1,)).item()
YJ = y[:, :, YJ_num, :]
YJc = y.clone()
YJc[:, :, YJ_num, :] = 0
weight = NP / (NP - 1)
RJc = self.reconstruct_measurements(YJc * weight, self.op)
output_recon_1 = self.model(RJc)
output_sino_1 = self.projector(output_recon_1)
AJ = output_sino_1[:, :, YJ_num, :]
batch_loss = ((AJ - YJ) ** 2).mean()
angle = random.uniform(0, 360)
rotated_output_recon_1 = rotate(output_recon_1, angle, interpolation="bilinear")
rotated_sinogram = self.projector(rotated_output_recon_1)
with torch.no_grad():
noise = self._noise(rotated_sinogram)
rotated_sinogram_noisy = rotated_sinogram + noise
rotated_noisy_image = self.reconstruct_measurements(
rotated_sinogram_noisy, self.op
)
output_recon_2 = self.model(rotated_noisy_image)
equivariance = ((output_recon_2 - rotated_output_recon_1) ** 2).mean()
batch_loss += self.solver_params.equivariance_strength * equivariance
return self._regularized(batch_loss, RJc, output_recon_1)
def _forward(self, data):
# validation and testing evaluate the reconstruction
return self.reconstruct(data)
[docs]
def reconstruct(self, y):
"""The model's reconstructions of the measurements ``y``: of the
reconstructions ``recon_fn`` makes of them, as in training."""
input_recon = self.reconstruct_measurements(y, self.op)
output_recon = self.model(input_recon)
return output_recon
[docs]
@staticmethod
def cite(cite_format="bib"):
"""Where Equivariance2Inverse comes from: its citation, BibTeX ("bib") or "MLA"."""
return citation(
cite_format,
bib="""
@article{schut2025equivariance2inverse,
title={Equivariance2Inverse: A Practical Self-Supervised CT Reconstruction Method Benchmarked on Real, Limited-Angle, and Blurred Data},
author={Schut, Dirk Elias and Graas, Adriaan and van Liere, Robert and van Leeuwen, Tristan},
journal={arXiv preprint arXiv:2510.23317},
year={2025}
}
""",
mla="""
Schut, Dirk Elias, Adriaan Graas, Robert van Liere, and Tristan van Leeuwen. "Equivariance2Inverse: A Practical Self-Supervised CT Reconstruction Method Benchmarked on Real, Limited-Angle, and Blurred Data." arXiv preprint arXiv:2510.23317 (2025).
""",
)