Overloaded mask_largest_object to include both object maps and binary masks as inputs

extensions/chaeo/process.py

 import numpy as np
-from skimage.measure import label, regionprops_table
+import skimage
 
+from model_server.process import is_mask
 
 def mask_largest_object(
         img: np.ndarray,
@@ -9,14 +10,15 @@ def mask_largest_object(
 ) -> np.ndarray:
     """
     Where more than one connected component is found in an image, return the largest object by area
-    :param img: (np.ndarray) containing object labels
+    :param img: (np.ndarray) containing object labels or binary mask
     :param max_allowed: raise an error if more than this number of objects is found
     :param verbose: print a message each time more than one object is found
     :return: np.ndarray of same size as img
     """
-    binary = img > 0
-    ob_id = label(binary)
-    # separate problem for mask and object maps
+    if is_mask(img): # assign object labels
+        ob_id = skimage.measure.label(img)
+    else:  # assume img is contains object labels
+        ob_id = img
 
     # import skimage
     # from pathlib import Path
@@ -28,9 +30,11 @@ def mask_largest_object(
     if num_obj > 1:
         if verbose:
             print(f'Found {num_obj} nonzero unique values in object map; keeping the one with the largest area')
-        pr = regionprops_table(ob_id, properties=['label', 'area'])
-        idx_max_area = pr['area'].argmax()
-        mask = ob_id == pr['label'][idx_max_area]
+        # pr = regionprops_table(ob_id, properties=['label', 'area'])
+        val, cts = np.unique(ob_id, return_counts=True)
+        mask = ob_id == val[1 + cts[1:].argmax()]
+        # idx_max_area = pr['area'].argmax()
+        # mask = ob_id == pr['label'][idx_max_area]
         return mask * img
     else:
         return img

extensions/chaeo/tests/test_process.py

@@ -5,17 +5,28 @@ import numpy as np
 from extensions.chaeo.process import mask_largest_object
 
 class TestMaskLargestObject(unittest.TestCase):
-     def test_mask_largest_object(self):
-         arr = np.zeros([5, 5])
-         arr[0:3, 0:3] = 2
-         arr[4, 2:5] = 4
-         masked = mask_largest_object(arr)
-         self.assertTrue(np.all(np.unique(masked) == [0, 2]))
-         self.assertTrue(np.all(masked[4:5, :] == 0))
-         self.assertTrue(np.all(masked[:, 4:5] == 0))
+    def test_mask_largest_touching_object(self):
+        arr = np.zeros([5, 5], dtype='uint8')
+        arr[0:3, 0:3] = 2
+        arr[3:, 2:] = 4
+        masked = mask_largest_object(arr)
+        self.assertTrue(np.all(np.unique(masked) == [0, 2]))
+        self.assertTrue(np.all(masked[4:5, 0:2] == 0))
+        self.assertTrue(np.all(masked[0:3, 3:5] == 0))
 
-     def test_no_change(self):
-        arr = np.zeros([5, 5])
+    def test_no_change(self):
+        arr = np.zeros([5, 5], dtype='uint8')
         arr[0:3, 0:3] = 2
         masked = mask_largest_object(arr)
         self.assertTrue(np.all(masked == arr))
+
+    def test_mask_multiple_objects_in_binary_maks(self):
+        arr = np.zeros([5, 5], dtype='uint8')
+        arr[0:3, 0:3] = 255
+        arr[4, 2:5] = 255
+        masked = mask_largest_object(arr)
+        print(np.unique(masked))
+        self.assertTrue(np.all(np.unique(masked) == [0, 255]))
+        self.assertTrue(np.all(masked[:, 3:5] == 0))
+        self.assertTrue(np.all(masked[3:5, :] == 0))
+