Update SegmentationExtractor

src/roiextractors/extraction_tools.py

@@ -31,6 +31,7 @@
 DtypeType = DTypeLike
 IntType = Union[int, np.integer]
 FloatType = float
+NoneType = type(None)
 
 
 def raise_multi_channel_or_depth_not_implemented(extractor_name: str):
@@ -244,18 +245,16 @@ def read_numpy_memmap_video(
     return video_memap
 
 
-def _pixel_mask_extractor(image_mask_, _roi_ids) -> list:
-    """Convert image mask to pixel mask.
+def _pixel_mask_extractor(image_masks: np.ndarray) -> list:
+    """Convert image masks to pixel masks.
 
     Pixel masks are an alternative data format for storage of image masks which relies on the sparsity of the images.
     The location and weight of each non-zero pixel is stored for each mask.
 
     Parameters
     ----------
-    image_mask_: numpy.ndarray
+    image_masks: numpy.ndarray
         Dense representation of the ROIs with shape (number_of_rows, number_of_columns, number_of_rois).
-    _roi_ids: list
-        List of roi ids with length number_of_rois.
 
     Returns
     -------
@@ -265,11 +264,11 @@ def _pixel_mask_extractor(image_mask_, _roi_ids) -> list:
         the pixel.
     """
     pixel_mask_list = []
-    for i, roiid in enumerate(_roi_ids):
-        image_mask = np.array(image_mask_[:, :, i])
-        _locs = np.where(image_mask > 0)
-        _pix_values = image_mask[image_mask > 0]
-        pixel_mask_list.append(np.vstack((_locs[0], _locs[1], _pix_values)).T)
+    for i in range(image_masks.shape[2]):
+        image_mask = image_masks[:, :, i]
+        locs = np.where(image_mask > 0)
+        pix_values = image_mask[image_mask > 0]
+        pixel_mask_list.append(np.vstack((locs[0], locs[1], pix_values)).T)
     return pixel_mask_list
 
 
@@ -684,3 +683,31 @@ def get_package(
         f"\nThe required package'{package_name}' is not installed!\n"
         f"To install this package, please run\n\n\t{installation_instructions}\n"
     )
+
+
+def get_default_roi_locations_from_image_masks(image_masks: np.ndarray) -> np.ndarray:
+    """Calculate the default ROI locations from given image masks.
+
+    This function takes a 3D numpy array of image masks and computes the median
+    coordinates of the maximum values in each 2D mask. The result is a 2D numpy
+    array where each column represents the (x, y) coordinates of the ROI for
+    each mask.
+
+    Parameters
+    ----------
+    image_masks : np.ndarray
+        A 3D numpy array of shape (height, width, num_rois) containing the image masks.
+
+    Returns
+    -------
+    np.ndarray
+        A 2D numpy array of shape (2, num_rois) where each column contains the
+        (x, y) coordinates of the ROI for each mask.
+    """
+    num_rois = image_masks.shape[2]
+    roi_locations = np.zeros([2, num_rois], dtype="int")
+    for i in range(num_rois):
+        image_mask = image_masks[:, :, i]
+        max_value_indices = np.where(image_mask == np.amax(image_mask))
+        roi_locations[:, i] = np.array([np.median(max_value_indices[0]), np.median(max_value_indices[1])]).T
+    return roi_locations