Skip to content

Commit

Permalink
fix implementation mistakes and add conjugate gradients solver (Proje…
Browse files Browse the repository at this point in the history
…ct-MONAI#7876)

Fixes Project-MONAI#6767 
### Description

Fixed two minor implementation differences between the official MATLAB
code and the MONAI implementation, to confirm also added a new test case
where two images and their confidence maps calculated by the official
code is added to tests and the MONAI implementation results are checked
against there results created by the official code.

Also fixing the issue: added the conjugate gradients solver option. Now
the users can utilize it to run the algorithm faster with a trade-off of
accuracy of the end result, a range of speed-ups can be achieved with
little to no quality loss by tweaking the parameters, the optimal
parameters between quality and speed in my experience is set as the
default parameters, namely 'cg_tol' and 'cg_maxiter'.

For the CG solver installing PyAMG (https://github.com/pyamg/pyamg) is a
requirement, this is because we use it to generate a preconditioner,
without it CG does not provide any speed-ups, even slows down the
algorithm. This part can be changed if the requirement is not ideal, yet
this was the best solution as far as my knowledge goes.

### Types of changes
<!--- Put an `x` in all the boxes that apply, and remove the not
applicable items -->
- [x] Non-breaking change (fix or new feature that would not break
existing functionality).
- [x] New tests added to cover the changes.
- [x] Integration tests passed locally by running `./runtests.sh -f -u
--net --coverage`.
- [x] Quick tests passed locally by running `./runtests.sh --quick
--unittests --disttests`.
- [x] In-line docstrings updated.
- [x] Documentation updated, tested `make html` command in the `docs/`
folder.

---------

Signed-off-by: MrGranddy <[email protected]>
Signed-off-by: Vahit Buğra YEŞİLKAYNAK <[email protected]>
Co-authored-by: ge85evz <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Eric Kerfoot <[email protected]>
Co-authored-by: YunLiu <[email protected]>
  • Loading branch information
5 people authored Jun 28, 2024
1 parent ac84a4e commit 06cbd70
Show file tree
Hide file tree
Showing 11 changed files with 531 additions and 455 deletions.
1 change: 1 addition & 0 deletions docs/requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -40,3 +40,4 @@ onnx>=1.13.0
onnxruntime; python_version <= '3.10'
zarr
huggingface_hub
pyamg>=5.0.0
2 changes: 1 addition & 1 deletion docs/source/installation.md
Original file line number Diff line number Diff line change
Expand Up @@ -258,6 +258,6 @@ Since MONAI v0.2.0, the extras syntax such as `pip install 'monai[nibabel]'` is
```

which correspond to `nibabel`, `scikit-image`,`scipy`, `pillow`, `tensorboard`,
`gdown`, `pytorch-ignite`, `torchvision`, `itk`, `tqdm`, `lmdb`, `psutil`, `cucim`, `openslide-python`, `pandas`, `einops`, `transformers`, `mlflow`, `clearml`, `matplotlib`, `tensorboardX`, `tifffile`, `imagecodecs`, `pyyaml`, `fire`, `jsonschema`, `ninja`, `pynrrd`, `pydicom`, `h5py`, `nni`, `optuna`, `onnx`, `onnxruntime`, `zarr`, `lpips`, `nvidia-ml-py`, and `huggingface_hub` respectively.
`gdown`, `pytorch-ignite`, `torchvision`, `itk`, `tqdm`, `lmdb`, `psutil`, `cucim`, `openslide-python`, `pandas`, `einops`, `transformers`, `mlflow`, `clearml`, `matplotlib`, `tensorboardX`, `tifffile`, `imagecodecs`, `pyyaml`, `fire`, `jsonschema`, `ninja`, `pynrrd`, `pydicom`, `h5py`, `nni`, `optuna`, `onnx`, `onnxruntime`, `zarr`, `lpips`, `nvidia-ml-py`, `huggingface_hub` and `pyamg` respectively.

- `pip install 'monai[all]'` installs all the optional dependencies.
45 changes: 38 additions & 7 deletions monai/data/ultrasound_confidence_map.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,9 @@
cv2, _ = optional_import("cv2")
csc_matrix, _ = optional_import("scipy.sparse", "1.7.1", min_version, "csc_matrix")
spsolve, _ = optional_import("scipy.sparse.linalg", "1.7.1", min_version, "spsolve")
cg, _ = optional_import("scipy.sparse.linalg", "1.7.1", min_version, "cg")
hilbert, _ = optional_import("scipy.signal", "1.7.1", min_version, "hilbert")
ruge_stuben_solver, _ = optional_import("pyamg", "5.0.0", min_version, "ruge_stuben_solver")


class UltrasoundConfidenceMap:
Expand All @@ -30,22 +32,43 @@ class UltrasoundConfidenceMap:
It generates a confidence map by setting source and sink points in the image and computing the probability
for random walks to reach the source for each pixel.
The official code is available at:
https://campar.in.tum.de/Main/AthanasiosKaramalisCode
Args:
alpha (float, optional): Alpha parameter. Defaults to 2.0.
beta (float, optional): Beta parameter. Defaults to 90.0.
gamma (float, optional): Gamma parameter. Defaults to 0.05.
mode (str, optional): 'RF' or 'B' mode data. Defaults to 'B'.
sink_mode (str, optional): Sink mode. Defaults to 'all'. If 'mask' is selected, a mask must be when calling
the transform. Can be 'all', 'mid', 'min', or 'mask'.
use_cg (bool, optional): Use Conjugate Gradient method for solving the linear system. Defaults to False.
cg_tol (float, optional): Tolerance for the Conjugate Gradient method. Defaults to 1e-6.
Will be used only if `use_cg` is True.
cg_maxiter (int, optional): Maximum number of iterations for the Conjugate Gradient method. Defaults to 200.
Will be used only if `use_cg` is True.
"""

def __init__(self, alpha: float = 2.0, beta: float = 90.0, gamma: float = 0.05, mode="B", sink_mode="all"):
def __init__(
self,
alpha: float = 2.0,
beta: float = 90.0,
gamma: float = 0.05,
mode="B",
sink_mode="all",
use_cg=False,
cg_tol=1e-6,
cg_maxiter=200,
):
# The hyperparameters for confidence map estimation
self.alpha = alpha
self.beta = beta
self.gamma = gamma
self.mode = mode
self.sink_mode = sink_mode
self.use_cg = use_cg
self.cg_tol = cg_tol
self.cg_maxiter = cg_maxiter

# The precision to use for all computations
self.eps = np.finfo("float64").eps
Expand Down Expand Up @@ -228,17 +251,18 @@ def confidence_laplacian(self, padded_index: NDArray, padded_image: NDArray, bet
s = self.normalize(s)

# Horizontal penalty
s[:vertical_end] += gamma
# s[vertical_end:diagonal_end] += gamma * np.sqrt(2) # --> In the paper it is sqrt(2)
# since the diagonal edges are longer yet does not exist in the original code
s[vertical_end:] += gamma
# Here there is a difference between the official MATLAB code and the paper
# on the edge penalty. We directly implement what the official code does.

# Normalize differences
s = self.normalize(s)

# Gaussian weighting function
s = -(
(np.exp(-beta * s, dtype="float64")) + 1.0e-6
) # --> This epsilon changes results drastically default: 1.e-6
(np.exp(-beta * s, dtype="float64")) + 1e-5
) # --> This epsilon changes results drastically default: 10e-6
# Please notice that it is not 1e-6, it is 10e-6 which is actually different.

# Create Laplacian, diagonal missing
lap = csc_matrix((s, (i, j)))
Expand All @@ -256,7 +280,14 @@ def confidence_laplacian(self, padded_index: NDArray, padded_image: NDArray, bet
return lap

def _solve_linear_system(self, lap, rhs):
x = spsolve(lap, rhs)

if self.use_cg:
lap_sparse = lap.tocsr()
ml = ruge_stuben_solver(lap_sparse, coarse_solver="pinv")
m = ml.aspreconditioner(cycle="V")
x, _ = cg(lap, rhs, tol=self.cg_tol, maxiter=self.cg_maxiter, M=m)
else:
x = spsolve(lap, rhs)

return x

Expand Down
27 changes: 25 additions & 2 deletions monai/transforms/intensity/array.py
Original file line number Diff line number Diff line change
Expand Up @@ -2789,21 +2789,42 @@ class UltrasoundConfidenceMapTransform(Transform):
It generates a confidence map by setting source and sink points in the image and computing the probability
for random walks to reach the source for each pixel.
The official code is available at:
https://campar.in.tum.de/Main/AthanasiosKaramalisCode
Args:
alpha (float, optional): Alpha parameter. Defaults to 2.0.
beta (float, optional): Beta parameter. Defaults to 90.0.
gamma (float, optional): Gamma parameter. Defaults to 0.05.
mode (str, optional): 'RF' or 'B' mode data. Defaults to 'B'.
sink_mode (str, optional): Sink mode. Defaults to 'all'. If 'mask' is selected, a mask must be when
calling the transform. Can be one of 'all', 'mid', 'min', 'mask'.
use_cg (bool, optional): Use Conjugate Gradient method for solving the linear system. Defaults to False.
cg_tol (float, optional): Tolerance for the Conjugate Gradient method. Defaults to 1e-6.
Will be used only if `use_cg` is True.
cg_maxiter (int, optional): Maximum number of iterations for the Conjugate Gradient method. Defaults to 200.
Will be used only if `use_cg` is True.
"""

def __init__(self, alpha: float = 2.0, beta: float = 90.0, gamma: float = 0.05, mode="B", sink_mode="all") -> None:
def __init__(
self,
alpha: float = 2.0,
beta: float = 90.0,
gamma: float = 0.05,
mode="B",
sink_mode="all",
use_cg=False,
cg_tol: float = 1.0e-6,
cg_maxiter: int = 200,
):
self.alpha = alpha
self.beta = beta
self.gamma = gamma
self.mode = mode
self.sink_mode = sink_mode
self.use_cg = use_cg
self.cg_tol = cg_tol
self.cg_maxiter = cg_maxiter

if self.mode not in ["B", "RF"]:
raise ValueError(f"Unknown mode: {self.mode}. Supported modes are 'B' and 'RF'.")
Expand All @@ -2813,7 +2834,9 @@ def __init__(self, alpha: float = 2.0, beta: float = 90.0, gamma: float = 0.05,
f"Unknown sink mode: {self.sink_mode}. Supported modes are 'all', 'mid', 'min' and 'mask'."
)

self._compute_conf_map = UltrasoundConfidenceMap(self.alpha, self.beta, self.gamma, self.mode, self.sink_mode)
self._compute_conf_map = UltrasoundConfidenceMap(
self.alpha, self.beta, self.gamma, self.mode, self.sink_mode, self.use_cg, self.cg_tol, self.cg_maxiter
)

def __call__(self, img: NdarrayOrTensor, mask: NdarrayOrTensor | None = None) -> NdarrayOrTensor:
"""Compute confidence map from an ultrasound image.
Expand Down
1 change: 1 addition & 0 deletions requirements-dev.txt
Original file line number Diff line number Diff line change
Expand Up @@ -57,3 +57,4 @@ zarr
lpips==0.1.4
nvidia-ml-py
huggingface_hub
pyamg>=5.0.0
3 changes: 3 additions & 0 deletions setup.cfg
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,7 @@ all =
lpips==0.1.4
nvidia-ml-py
huggingface_hub
pyamg>=5.0.0
nibabel =
nibabel
ninja =
Expand Down Expand Up @@ -162,6 +163,8 @@ pynvml =
# MetricsReloaded @ git+https://github.com/Project-MONAI/MetricsReloaded@monai-support#egg=MetricsReloaded
huggingface_hub =
huggingface_hub
pyamg =
pyamg>=5.0.0

[flake8]
select = B,C,E,F,N,P,T4,W,B9
Expand Down
Loading

0 comments on commit 06cbd70

Please sign in to comment.