fix implementation mistakes and add conjugate gradients solver

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:
@@ -37,15 +39,31 @@ class UltrasoundConfidenceMap:
         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.
+        cg_maxiter (int, optional): Maximum number of iterations for the Conjugate Gradient method. Defaults to 200.
     """
 
-    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 +246,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")) + 10e-6
+        )  # --> 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 +275,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, rtol=self.cg_tol, maxiter=self.cg_maxiter, M=m)
+        else:
+            x = spsolve(lap, rhs)
 
         return x
 

monai/transforms/intensity/array.py

@@ -2796,14 +2796,30 @@ class UltrasoundConfidenceMapTransform(Transform):
         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.
+        cg_maxiter (int, optional): Maximum number of iterations for the Conjugate Gradient method. Defaults to 200.
     """
 
-    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 +2829,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

@@ -26,7 +26,7 @@ mypy>=1.5.0
 ninja
 torchvision
 psutil
-cucim-cu12; platform_system == "Linux" and python_version >= "3.9" and python_version <= "3.10"
+#cucim-cu12; platform_system == "Linux" and python_version >= "3.9" and python_version <= "3.10"
 openslide-python
 imagecodecs; platform_system == "Linux" or platform_system == "Darwin"
 tifffile; platform_system == "Linux" or platform_system == "Darwin"
@@ -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