fix implementation mistakes and add conjugate gradients solver (#7876)

Fixes #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]>

docs/requirements.txt

@@ -40,3 +40,4 @@ onnx>=1.13.0
 onnxruntime; python_version <= '3.10'
 zarr
 huggingface_hub
+pyamg>=5.0.0

docs/source/installation.md

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

monai/data/ultrasound_confidence_map.py

@@ -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:
@@ -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
@@ -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)))
@@ -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
 

monai/transforms/intensity/array.py

@@ -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'.")
@@ -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.

requirements-dev.txt

@@ -57,3 +57,4 @@ zarr
 lpips==0.1.4
 nvidia-ml-py
 huggingface_hub
+pyamg>=5.0.0

setup.cfg

@@ -84,6 +84,7 @@ all =
     lpips==0.1.4
     nvidia-ml-py
     huggingface_hub
+    pyamg>=5.0.0
 nibabel =
     nibabel
 ninja =
@@ -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