From c0c666915fbb3095a8aa23b0e4f58ea14ad756af Mon Sep 17 00:00:00 2001
From: Nathan Painchaud <23144457+nathanpainchaud@users.noreply.github.com>
Date: Wed, 3 Aug 2022 23:28:46 +0200
Subject: [PATCH] Improve `EchoMeasure` to avoid crash when candidate markers
 are aligned vertically (#76)

---
 vital/utils/image/us/measure.py | 21 +++++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/vital/utils/image/us/measure.py b/vital/utils/image/us/measure.py
index 8f0c49b7..83f18705 100644
--- a/vital/utils/image/us/measure.py
+++ b/vital/utils/image/us/measure.py
@@ -43,7 +43,7 @@ def _lv_base(
         others = ~(left_ventricle + myocardium)
         dilated_myocardium = ndimage.binary_dilation(myocardium, structure=EchoMeasure._square_3x3_connectivity)
         dilated_others = ndimage.binary_dilation(others, structure=EchoMeasure._square_3x3_connectivity)
-        y_coords, x_coords = np.where(left_ventricle * dilated_myocardium * dilated_others)
+        y_coords, x_coords = np.nonzero(left_ventricle * dilated_myocardium * dilated_others)
 
         if (num_markers := len(y_coords)) < 2:
             logger.warning(
@@ -52,12 +52,21 @@ def _lv_base(
             )
             return np.nan, np.nan
 
-        left_half_mask = x_coords < x_coords.mean()
-        bottom_left_point_idx = y_coords[left_half_mask].argmax()
-        bottom_right_point_idx = y_coords[~left_half_mask].argmax()
+        if np.all(x_coords == x_coords.mean()):
+            # Edge case where the base points are aligned vertically
+            # Divide frontier into bottom and top halves.
+            coord_mask = y_coords > y_coords.mean()
+            left_point_idx = y_coords[coord_mask].argmin()
+            right_point_idx = y_coords[~coord_mask].argmax()
+        else:
+            # Normal case where there is a clear divide between left and right markers at the base
+            # Divide frontier into left and right halves.
+            coord_mask = x_coords < x_coords.mean()
+            left_point_idx = y_coords[coord_mask].argmax()
+            right_point_idx = y_coords[~coord_mask].argmax()
         return (
-            (y_coords[left_half_mask][bottom_left_point_idx], x_coords[left_half_mask][bottom_left_point_idx]),
-            (y_coords[~left_half_mask][bottom_right_point_idx], x_coords[~left_half_mask][bottom_right_point_idx]),
+            (y_coords[coord_mask][left_point_idx], x_coords[coord_mask][left_point_idx]),
+            (y_coords[~coord_mask][right_point_idx], x_coords[~coord_mask][right_point_idx]),
         )
 
     @staticmethod