Repository: Project-MONAI/MONAI

Issue:
faster solver for computing confidence maps
follow up of #6709
             
 > The issue though is the process for installing Octave and oct2py is more complex than other things we rely on, requiring a separate install and a correctly configured PATH variable. This is also an issue for our CICD system since we do need to test this transform. The library you're using, oct2py, also seems like a rather small project whose long term support is not assured. I'd still recommend we look into some other solver (something wrapping eigen?) to avoid this dependency.

Hello, for now I've removed the octave dependency, now the code only uses scipy solver.
As far as my knowledge goes, there is no decomposition or iterative method that would make
the algorithm faster while preserving the results. I am also sharing some metrics for the comparison
of octave and scipy. I believe SciPy is not all bad, considering they also would optimize their sparse solver
over time.

![image](https://github.com/Project-MONAI/MONAI/assets/17518820/e5d6a33c-397f-4df2-85d5-48e7258896dd)

For an image with 118677 pixels (around 350 x 350), using i7-13700K with 10 runs each:

```
Scipy mean time: 5.763958525657654 +- 0.2072540601999276
Octave mean time: 0.7772286891937256 +- 0.018855939984020915
```

_Originally posted by @MrGranddy in https://github.com/Project-MONAI/MONAI/issues/6709#issuecomment-1641391810_
            


Discussion:
I actually achieved satisfactory results with "Conjugate Gradients" method recently, it is of course not deterministicly solving the system so there will be some minor differences compared to the original one, yet I believe it is close enough and even not so, the user always can trade between accuracy and speed according to their need if the parameters controlling that is also given. I will soon push the new version.

Basically the 2 parameters controlling the algorithm are `max_iter`: maximum number of iterations taken by the algorithm,
`tol` floating point numerical error to chech convergence, lower `tol` means better accuracy and slower operation. `max_iter` can also be used to put an upper limit for time required.

If you have any comments now, I can follow them before the initial commit, for example how to implement these "2-mode" behaviour one `deterministic` and the other `iterative` best in MONAI. You can of course also wait until initial commit if it makes more sense.

Here is a showcase of the new algorithm:

![image](https://github.com/Project-MONAI/MONAI/assets/17518820/8fbd20b1-44aa-4ce5-ba10-55e97449eab3)


Resolving patch:
diff --git a/monai/data/ultrasound_confidence_map.py b/monai/data/ultrasound_confidence_map.py
--- a/monai/data/ultrasound_confidence_map.py
+++ b/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,6 +32,9 @@ 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.
@@ -37,15 +42,33 @@ 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.
+            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
 
diff --git a/monai/transforms/intensity/array.py b/monai/transforms/intensity/array.py
--- a/monai/transforms/intensity/array.py
+++ b/monai/transforms/intensity/array.py
@@ -2789,6 +2789,9 @@ 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.
@@ -2796,14 +2799,32 @@ 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.
+            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.


Test patch:
diff --git a/tests/test_ultrasound_confidence_map_transform.py b/tests/test_ultrasound_confidence_map_transform.py
--- a/tests/test_ultrasound_confidence_map_transform.py
+++ b/tests/test_ultrasound_confidence_map_transform.py
@@ -11,11 +11,13 @@
 
 from __future__ import annotations
 
+import os
 import unittest
 
 import numpy as np
 import torch
 from parameterized import parameterized
+from PIL import Image
 
 from monai.transforms import UltrasoundConfidenceMapTransform
 from tests.utils import assert_allclose
@@ -32,7 +34,8 @@
         [1, 2, 3, 32, 33, 34, 35, 1, 2, 3],
         [1, 2, 3, 36, 37, 38, 39, 1, 2, 3],
         [1, 2, 3, 40, 41, 42, 43, 1, 2, 3],
-    ]
+    ],
+    dtype=np.float32,
 )
 
 TEST_MASK = np.array(
@@ -47,474 +50,435 @@
         [1, 1, 1, 0, 0, 0, 1, 1, 1, 0],
         [1, 1, 1, 0, 0, 0, 1, 1, 1, 0],
         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
-    ]
+    ],
+    dtype=np.float32,
 )
 
 SINK_ALL_OUTPUT = np.array(
     [
         [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
         [
-            0.97514489,
-            0.96762971,
-            0.96164186,
-            0.95463443,
-            0.9941512,
-            0.99023054,
-            0.98559401,
-            0.98230057,
-            0.96601224,
-            0.95119599,
-        ],
-        [
-            0.92960533,
-            0.92638451,
-            0.9056675,
-            0.9487176,
-            0.9546961,
-            0.96165853,
-            0.96172303,
-            0.92686401,
-            0.92122613,
-            0.89957239,
-        ],
-        [
-            0.86490963,
-            0.85723665,
-            0.83798141,
-            0.90816201,
-            0.90816097,
-            0.90815301,
-            0.9081427,
-            0.85933627,
-            0.85146935,
-            0.82948586,
-        ],
-        [
-            0.77430346,
-            0.76731372,
-            0.74372311,
-            0.89128774,
-            0.89126885,
-            0.89125066,
-            0.89123521,
-            0.76858589,
-            0.76106647,
-            0.73807776,
-        ],
-        [
-            0.66098109,
-            0.65327697,
-            0.63090644,
-            0.33086588,
-            0.3308383,
-            0.33081937,
-            0.33080718,
-            0.6557468,
-            0.64825099,
-            0.62593375,
-        ],
-        [
-            0.52526945,
-            0.51832586,
-            0.49709412,
-            0.25985059,
-            0.25981009,
-            0.25977729,
-            0.25975222,
-            0.52118958,
-            0.51426328,
-            0.49323164,
-        ],
-        [
-            0.3697845,
-            0.36318971,
-            0.34424661,
-            0.17386804,
-            0.17382046,
-            0.17377993,
-            0.17374668,
-            0.36689317,
-            0.36036096,
-            0.3415582,
-        ],
-        [
-            0.19546374,
-            0.1909659,
-            0.17319999,
-            0.08423318,
-            0.08417993,
-            0.08413242,
-            0.08409104,
-            0.19393909,
-            0.18947485,
-            0.17185031,
+            0.8884930952884654,
+            0.8626656901726876,
+            0.8301161870669913,
+            0.9757179300830185,
+            0.9989819637626414,
+            0.9994717624885747,
+            0.9954377526794013,
+            0.8898638133944221,
+            0.862604343021387,
+            0.8277862494812598,
+        ],
+        [
+            0.7765718877433174,
+            0.7363731552518268,
+            0.6871875923653379,
+            0.9753673327387775,
+            0.9893175316399789,
+            0.9944181334242039,
+            0.9936979128319371,
+            0.7778001700035326,
+            0.7362622619974832,
+            0.6848377775329241,
+        ],
+        [
+            0.6648416226360719,
+            0.6178079903692397,
+            0.5630152545966568,
+            0.8278402502498404,
+            0.82790391019578,
+            0.8289702087149963,
+            0.8286730258710652,
+            0.6658773633169731,
+            0.6176836507071695,
+            0.5609165245633834,
+        ],
+        [
+            0.5534420483956817,
+            0.5055401989946189,
+            0.451865872383879,
+            0.7541423053657541,
+            0.7544115886347456,
+            0.7536884376055174,
+            0.7524927915364896,
+            0.5542943466824017,
+            0.505422678400297,
+            0.4502051549732117,
+        ],
+        [
+            0.4423657561928356,
+            0.398221575954319,
+            0.35030055029978124,
+            0.4793202144786371,
+            0.48057175662074125,
+            0.4812057229564038,
+            0.48111949176149327,
+            0.44304092606050766,
+            0.39812149713417405,
+            0.34902458531143377,
+        ],
+        [
+            0.3315561576450342,
+            0.29476346732036784,
+            0.2558303772864961,
+            0.35090405668257535,
+            0.3515225984307705,
+            0.35176548159366317,
+            0.3516979775419521,
+            0.33205839061494885,
+            0.2946859567272435,
+            0.2549042599220772,
+        ],
+        [
+            0.22094175240967673,
+            0.19431840633358133,
+            0.16672448058324435,
+            0.22716195845848167,
+            0.22761996456848282,
+            0.22782525614780919,
+            0.22781876632199002,
+            0.22127471252104777,
+            0.19426593309729956,
+            0.16612306610996525,
+        ],
+        [
+            0.11044782531624744,
+            0.09623229814933323,
+            0.08174664901235043,
+            0.11081911718888311,
+            0.11102310514207447,
+            0.1111041051969924,
+            0.11108329076967229,
+            0.11061376973431204,
+            0.09620592927336903,
+            0.08145227209865454,
         ],
         [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
-    ]
+    ],
+    dtype=np.float32,
 )
 
 SINK_MID_OUTPUT = np.array(
     [
+        [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
         [
-            1.00000000e00,
-            1.00000000e00,
-            1.00000000e00,
-            1.00000000e00,
-            1.00000000e00,
-            1.00000000e00,
-            1.00000000e00,
-            1.00000000e00,
-            1.00000000e00,
-            1.00000000e00,
-        ],
-        [
-            9.99996103e-01,
-            9.99994823e-01,
-            9.99993550e-01,
-            9.99930863e-01,
-            9.99990782e-01,
-            9.99984683e-01,
-            9.99979000e-01,
-            9.99997804e-01,
-            9.99995985e-01,
-            9.99994325e-01,
-        ],
-        [
-            9.99989344e-01,
-            9.99988600e-01,
-            9.99984099e-01,
-            9.99930123e-01,
-            9.99926598e-01,
-            9.99824297e-01,
-            9.99815032e-01,
-            9.99991228e-01,
-            9.99990881e-01,
-            9.99988462e-01,
-        ],
-        [
-            9.99980787e-01,
-            9.99979264e-01,
-            9.99975828e-01,
-            9.59669286e-01,
-            9.59664779e-01,
-            9.59656566e-01,
-            9.59648332e-01,
-            9.99983882e-01,
-            9.99983038e-01,
-            9.99980732e-01,
-        ],
-        [
-            9.99970181e-01,
-            9.99969032e-01,
-            9.99965730e-01,
-            9.45197806e-01,
-            9.45179593e-01,
-            9.45163629e-01,
-            9.45151458e-01,
-            9.99973352e-01,
-            9.99973254e-01,
-            9.99971098e-01,
-        ],
-        [
-            9.99958608e-01,
-            9.99957307e-01,
-            9.99953444e-01,
-            4.24743523e-01,
-            4.24713305e-01,
-            4.24694646e-01,
-            4.24685271e-01,
-            9.99960948e-01,
-            9.99961829e-01,
-            9.99960347e-01,
-        ],
-        [
-            9.99946675e-01,
-            9.99945139e-01,
-            9.99940312e-01,
-            3.51353224e-01,
-            3.51304003e-01,
-            3.51268260e-01,
-            3.51245366e-01,
-            9.99947688e-01,
-            9.99950165e-01,
-            9.99949512e-01,
-        ],
-        [
-            9.99935877e-01,
-            9.99934088e-01,
-            9.99928982e-01,
-            2.51197134e-01,
-            2.51130273e-01,
-            2.51080014e-01,
-            2.51045852e-01,
-            9.99936187e-01,
-            9.99939716e-01,
-            9.99940022e-01,
-        ],
-        [
-            9.99927846e-01,
-            9.99925911e-01,
-            9.99920188e-01,
-            1.31550973e-01,
-            1.31462736e-01,
-            1.31394558e-01,
-            1.31346069e-01,
-            9.99927275e-01,
-            9.99932142e-01,
-            9.99933313e-01,
-        ],
-        [
-            9.99924204e-01,
-            9.99922004e-01,
-            9.99915767e-01,
-            3.04861147e-04,
-            1.95998056e-04,
-            0.00000000e00,
-            2.05182682e-05,
-            9.99923115e-01,
-            9.99928835e-01,
-            9.99930535e-01,
-        ],
-    ]
+            0.9999957448889315,
+            0.9999781044114231,
+            0.9999142422442185,
+            0.999853253199584,
+            0.9999918403054282,
+            0.9999874855193227,
+            0.9999513619364747,
+            0.9999589247003497,
+            0.9999861765528631,
+            0.9999939213967494,
+        ],
+        [
+            0.9999918011366045,
+            0.9999588498417253,
+            0.9998388659316617,
+            0.9998496524281603,
+            0.9999154673258592,
+            0.9997827845182361,
+            0.9998160234579786,
+            0.9999163964511287,
+            0.9999743435786168,
+            0.9999894752861168,
+        ],
+        [
+            0.9999883847481621,
+            0.9999427334014465,
+            0.9997703972600652,
+            0.9853967608835997,
+            0.9852517829915376,
+            0.9853308520519438,
+            0.9854102394414211,
+            0.9998728503298413,
+            0.9999642585978225,
+            0.999986204909933,
+        ],
+        [
+            0.999985544721449,
+            0.9999296195017368,
+            0.9997066149628903,
+            0.9753803016111353,
+            0.9750688049429371,
+            0.9749211929217173,
+            0.9750052047129354,
+            0.9998284130289159,
+            0.9999558481338295,
+            0.9999837966320273,
+        ],
+        [
+            0.9999832723447848,
+            0.9999192263814408,
+            0.9996472692076177,
+            0.90541293509353,
+            0.9049945536526819,
+            0.9051142437853055,
+            0.9057005861296792,
+            0.9997839348839027,
+            0.9999490318922627,
+            0.9999820419085812,
+        ],
+        [
+            0.9999815409510937,
+            0.9999113168889934,
+            0.9995930143319085,
+            0.8370025145062345,
+            0.8358345435164332,
+            0.8358231468627223,
+            0.8369430449157075,
+            0.9997408260265034,
+            0.9999437526409107,
+            0.9999808010740554,
+        ],
+        [
+            0.9999803198262347,
+            0.9999057164296593,
+            0.9995461103528891,
+            0.7047260555380003,
+            0.7023346743490383,
+            0.7022946969603594,
+            0.7045662738042475,
+            0.9997017258131392,
+            0.9999399744001316,
+            0.9999799785302944,
+        ],
+        [
+            0.9999795785255197,
+            0.9999022923125928,
+            0.999510772973329,
+            0.46283993237260707,
+            0.4577365087549323,
+            0.4571888733219068,
+            0.4614967878524538,
+            0.9996710272733927,
+            0.9999376682163403,
+            0.9999795067125865,
+        ],
+        [
+            0.9999792877553907,
+            0.9999009179811408,
+            0.9994950057121632,
+            0.05049460567213739,
+            0.030946131978013824,
+            0.0,
+            0.019224121648385283,
+            0.9996568912408903,
+            0.9999367861122628,
+            0.9999793358521326,
+        ],
+    ],
+    dtype=np.float32,
 )
 
 SINK_MIN_OUTPUT = np.array(
     [
         [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
         [
-            0.99997545,
-            0.99996582,
-            0.99995245,
-            0.99856594,
-            0.99898314,
-            0.99777223,
-            0.99394423,
-            0.98588113,
-            0.97283215,
-            0.96096504,
-        ],
-        [
-            0.99993872,
-            0.99993034,
-            0.9998832,
-            0.9986147,
-            0.99848741,
-            0.9972981,
-            0.99723719,
-            0.94157173,
-            0.9369832,
-            0.91964243,
-        ],
-        [
-            0.99990802,
-            0.99989475,
-            0.99986873,
-            0.98610197,
-            0.98610047,
-            0.98609749,
-            0.98609423,
-            0.88741275,
-            0.88112911,
-            0.86349156,
-        ],
-        [
-            0.99988924,
-            0.99988509,
-            0.99988698,
-            0.98234089,
-            0.98233591,
-            0.98233065,
-            0.98232562,
-            0.81475172,
-            0.80865978,
-            0.79033138,
-        ],
-        [
-            0.99988418,
-            0.99988484,
-            0.99988323,
-            0.86796555,
-            0.86795874,
-            0.86795283,
-            0.86794756,
-            0.72418193,
-            0.71847704,
-            0.70022037,
-        ],
-        [
-            0.99988241,
-            0.99988184,
-            0.99988103,
-            0.85528225,
-            0.85527303,
-            0.85526389,
-            0.85525499,
-            0.61716519,
-            0.61026209,
-            0.59503671,
-        ],
-        [
-            0.99988015,
-            0.99987985,
-            0.99987875,
-            0.84258114,
-            0.84257121,
-            0.84256042,
-            0.84254897,
-            0.48997924,
-            0.49083978,
-            0.46891561,
-        ],
-        [
-            0.99987865,
-            0.99987827,
-            0.9998772,
-            0.83279589,
-            0.83278624,
-            0.83277384,
-            0.83275897,
-            0.36345545,
-            0.33690244,
-            0.35696828,
-        ],
-        [
-            0.99987796,
-            0.99987756,
-            0.99987643,
-            0.82873223,
-            0.82872648,
-            0.82871803,
-            0.82870711,
-            0.0,
-            0.26106012,
-            0.29978657,
-        ],
-    ]
+            0.9999961997987318,
+            0.9999801752476248,
+            0.9999185667341594,
+            0.9993115972922259,
+            0.9999536433504382,
+            0.9997590064584757,
+            0.9963282396026231,
+            0.9020645423682648,
+            0.965641014946897,
+            0.9847003633599846,
+        ],
+        [
+            0.9999926824858815,
+            0.9999628275604145,
+            0.9998472915971415,
+            0.9992953054409239,
+            0.9995550237000549,
+            0.9972853256638443,
+            0.9958871482234863,
+            0.8006505271617617,
+            0.9360757301263053,
+            0.9734843475613124,
+        ],
+        [
+            0.9999896427490426,
+            0.9999484707116104,
+            0.9997841142091455,
+            0.9321779021295554,
+            0.9308591506422442,
+            0.9299937642438358,
+            0.9286536283468563,
+            0.6964658886602826,
+            0.9106656689679997,
+            0.9652109119709528,
+        ],
+        [
+            0.9999871227708508,
+            0.9999369646510842,
+            0.9997276125796202,
+            0.9006206490361908,
+            0.8987968702587018,
+            0.8965696900664386,
+            0.8941507574801211,
+            0.5892568658180841,
+            0.8892240419729905,
+            0.9590996257620853,
+        ],
+        [
+            0.9999851119906539,
+            0.9999280075234918,
+            0.9996788394671484,
+            0.778755271203017,
+            0.7763917808258874,
+            0.7737517385551721,
+            0.7707980517990098,
+            0.4788014936236403,
+            0.8715671104783401,
+            0.954632732759503,
+        ],
+        [
+            0.9999835837292402,
+            0.999921323618806,
+            0.9996389455307461,
+            0.7222961578407286,
+            0.7186158832946955,
+            0.7146983167265393,
+            0.7105768254632475,
+            0.3648911004360315,
+            0.8575943501305144,
+            0.9514642802768379,
+        ],
+        [
+            0.9999825081019064,
+            0.999916683268467,
+            0.9996093996776352,
+            0.6713490686473397,
+            0.6664914636518112,
+            0.6613110504728309,
+            0.6558325489984669,
+            0.247299682539502,
+            0.8473037957967624,
+            0.9493580587294981,
+        ],
+        [
+            0.999981856118739,
+            0.9999138938063622,
+            0.9995907248497593,
+            0.6331535096751639,
+            0.6271637176135582,
+            0.6206687804556549,
+            0.6136262027168252,
+            0.12576864809108962,
+            0.8407892431959736,
+            0.9481472656653798,
+        ],
+        [
+            0.9999816006081851,
+            0.9999127861527936,
+            0.9995832399159849,
+            0.6133274396648696,
+            0.6086364734302403,
+            0.6034602717119345,
+            0.5978473214165134,
+            0.0,
+            0.8382338778894218,
+            0.9477082231321966,
+        ],
+    ],
+    dtype=np.float32,
 )
 
 SINK_MASK_OUTPUT = np.array(
     [
+        [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
+        [0.0, 0.0, 0.0, 0.9047934405899283, 0.9936046284605553, 0.9448690902377527, 0.0, 0.0, 0.0, 0.8363773255131761],
+        [0.0, 0.0, 0.0, 0.90375200446097, 0.9434594475474036, 0.4716831449516178, 0.0, 0.0, 0.0, 0.7364197333910302],
+        [
+            0.0,
+            0.0,
+            0.0,
+            0.09080438801405301,
+            0.06774182873204163,
+            0.038207095016625024,
+            0.0,
+            0.0,
+            0.0,
+            0.6745641479264269,
+        ],
+        [
+            0.0,
+            0.0,
+            0.0,
+            0.01731082802870267,
+            0.013540929458217351,
+            0.007321202161532623,
+            0.0,
+            0.0,
+            0.0,
+            0.6341231654271253,
+        ],
+        [
+            0.0,
+            0.0,
+            0.0,
+            0.0006444251665178544,
+            0.0005397129128756325,
+            0.0003048384803626333,
+            0.0,
+            0.0,
+            0.0,
+            0.6070178708536365,
+        ],
+        [
+            0.0,
+            0.0,
+            0.0,
+            5.406078586212675e-05,
+            4.416783924970537e-05,
+            2.4597362039020103e-05,
+            0.0,
+            0.0,
+            0.0,
+            0.5889413683184284,
+        ],
         [
-            1.00000000e00,
-            1.00000000e00,
-            1.00000000e00,
-            1.00000000e00,
-            1.00000000e00,
-            1.00000000e00,
-            1.00000000e00,
-            1.00000000e00,
-            1.00000000e00,
-            1.00000000e00,
-        ],
-        [
-            0.00000000e00,
-            0.00000000e00,
-            0.00000000e00,
-            2.86416400e-01,
-            7.93271181e-01,
-            5.81341234e-01,
-            0.00000000e00,
-            0.00000000e00,
-            0.00000000e00,
-            1.98395623e-01,
-        ],
-        [
-            0.00000000e00,
-            0.00000000e00,
-            0.00000000e00,
-            2.66733297e-01,
-            2.80741490e-01,
-            4.14078784e-02,
-            0.00000000e00,
-            0.00000000e00,
-            0.00000000e00,
-            7.91676486e-04,
-        ],
-        [
-            0.00000000e00,
-            0.00000000e00,
-            0.00000000e00,
-            1.86244537e-04,
-            1.53413401e-04,
-            7.85806495e-05,
-            0.00000000e00,
-            0.00000000e00,
-            0.00000000e00,
-            5.09797387e-06,
-        ],
-        [
-            0.00000000e00,
-            0.00000000e00,
-            0.00000000e00,
-            9.62904581e-07,
-            7.23946225e-07,
-            3.68824440e-07,
-            0.00000000e00,
-            0.00000000e00,
-            0.00000000e00,
-            4.79525316e-08,
-        ],
-        [
-            0.00000000e00,
-            0.00000000e00,
-            0.00000000e00,
-            1.50939343e-10,
-            1.17724874e-10,
-            6.21760843e-11,
-            0.00000000e00,
-            0.00000000e00,
-            0.00000000e00,
-            6.08922784e-10,
-        ],
-        [
-            0.00000000e00,
-            0.00000000e00,
-            0.00000000e00,
-            2.57593754e-13,
-            1.94066716e-13,
-            9.83784370e-14,
-            0.00000000e00,
-            0.00000000e00,
-            0.00000000e00,
-            9.80828665e-12,
-        ],
-        [
-            0.00000000e00,
-            0.00000000e00,
-            0.00000000e00,
-            4.22323494e-16,
-            3.17556633e-16,
-            1.60789400e-16,
-            0.00000000e00,
-            0.00000000e00,
-            0.00000000e00,
-            1.90789819e-13,
-        ],
-        [
-            0.00000000e00,
-            0.00000000e00,
-            0.00000000e00,
-            7.72677888e-19,
-            5.83029424e-19,
-            2.95946659e-19,
-            0.00000000e00,
-            0.00000000e00,
-            0.00000000e00,
-            4.97038275e-15,
-        ],
-        [
-            2.71345908e-24,
-            5.92006757e-24,
-            2.25580089e-23,
-            3.82601970e-18,
-            3.82835349e-18,
-            3.83302158e-18,
-            3.84002606e-18,
-            8.40760586e-16,
-            1.83433696e-15,
-            1.11629633e-15,
-        ],
-    ]
+            0.0,
+            0.0,
+            0.0,
+            4.39259327223233e-06,
+            3.6050656774754658e-06,
+            2.0127120155893425e-06,
+            0.0,
+            0.0,
+            0.0,
+            0.5774279920364456,
+        ],
+        [
+            0.0,
+            0.0,
+            0.0,
+            4.0740501726718113e-07,
+            3.374875487404489e-07,
+            1.9113630985667455e-07,
+            0.0,
+            0.0,
+            0.0,
+            0.5709897726747111,
+        ],
+        [
+            3.2266922388030425e-17,
+            1.801110982679718e-14,
+            9.325899448306927e-12,
+            3.913608442133728e-07,
+            3.9581822403393465e-07,
+            4.02383505118481e-07,
+            4.14820241328287e-07,
+            4.281640797396309e-06,
+            0.0023900192231620593,
+            0.5686882523793125,
+        ],
+    ],
+    dtype=np.float32,
 )
 
 
@@ -527,6 +491,21 @@ def setUp(self):
         self.input_img_torch = torch.from_numpy(TEST_INPUT).unsqueeze(0)  # mock image (torch tensor)
         self.input_mask_torch = torch.from_numpy(TEST_MASK).unsqueeze(0)  # mock mask (torch tensor)
 
+        self.real_input_img_paths = [
+            os.path.join(os.path.dirname(__file__), "testing_data", "ultrasound_confidence_map", "neck_input.png"),
+            os.path.join(os.path.dirname(__file__), "testing_data", "ultrasound_confidence_map", "femur_input.png"),
+        ]
+
+        self.real_result_npy_paths = [
+            os.path.join(os.path.dirname(__file__), "testing_data", "ultrasound_confidence_map", "neck_result.npy"),
+            os.path.join(os.path.dirname(__file__), "testing_data", "ultrasound_confidence_map", "femur_result.npy"),
+        ]
+
+        self.real_input_paramaters = [
+            {"alpha": 2.0, "beta": 90, "gamma": 0.03},
+            {"alpha": 2.0, "beta": 90, "gamma": 0.06},
+        ]
+
     def test_parameters(self):
         # Unknown mode
         with self.assertRaises(ValueError):
@@ -683,6 +662,44 @@ def test_func(self):
         output = transform(self.input_img_torch, self.input_mask_torch)
         assert_allclose(output, torch.tensor(SINK_MASK_OUTPUT), rtol=1e-4, atol=1e-4)
 
+    def test_against_official_code(self):
+        # This test is to compare the output of the transform with the official code
+        # The official code is available at:
+        # https://campar.in.tum.de/Main/AthanasiosKaramalisCode
+
+        for input_img_path, result_npy_path, params in zip(
+            self.real_input_img_paths, self.real_result_npy_paths, self.real_input_paramaters
+        ):
+            input_img = np.array(Image.open(input_img_path))
+            input_img = np.expand_dims(input_img, axis=0)
+
+            result_img = np.load(result_npy_path)
+
+            transform = UltrasoundConfidenceMapTransform(sink_mode="all", **params)
+            output = transform(input_img)
+
+            assert_allclose(output, result_img, rtol=1e-4, atol=1e-4)
+
+    def test_against_official_code_using_cg(self):
+        # This test is to compare the output of the transform with the official code
+        # The official code is available at:
+        # https://campar.in.tum.de/Main/AthanasiosKaramalisCode
+
+        for input_img_path, result_npy_path, params in zip(
+            self.real_input_img_paths, self.real_result_npy_paths, self.real_input_paramaters
+        ):
+            input_img = np.array(Image.open(input_img_path))
+            input_img = np.expand_dims(input_img, axis=0)
+
+            result_img = np.load(result_npy_path)
+
+            transform = UltrasoundConfidenceMapTransform(
+                sink_mode="all", use_cg=True, cg_tol=1.0e-6, cg_maxiter=300, **params
+            )
+            output = transform(input_img)
+
+            assert_allclose(output, result_img, rtol=1e-2, atol=1e-2)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/tests/testing_data/ultrasound_confidence_map/femur_input.png b/tests/testing_data/ultrasound_confidence_map/femur_input.png
new file mode 100644
Binary files /dev/null and b/tests/testing_data/ultrasound_confidence_map/femur_input.png differ
diff --git a/tests/testing_data/ultrasound_confidence_map/femur_result.npy b/tests/testing_data/ultrasound_confidence_map/femur_result.npy
new file mode 100644
Binary files /dev/null and b/tests/testing_data/ultrasound_confidence_map/femur_result.npy differ
diff --git a/tests/testing_data/ultrasound_confidence_map/neck_input.png b/tests/testing_data/ultrasound_confidence_map/neck_input.png
new file mode 100644
Binary files /dev/null and b/tests/testing_data/ultrasound_confidence_map/neck_input.png differ
diff --git a/tests/testing_data/ultrasound_confidence_map/neck_result.npy b/tests/testing_data/ultrasound_confidence_map/neck_result.npy
new file mode 100644
Binary files /dev/null and b/tests/testing_data/ultrasound_confidence_map/neck_result.npy differ