Merge pull request #192 from fractal-analytics-platform/132_bounding_box

ref #132, add preliminary implementation of bounding box

Author: tcompa

fractal_tasks_core/cellpose_segmentation.py

@@ -26,15 +26,23 @@
 import anndata as ad
 import dask.array as da
 import numpy as np
+import pandas as pd
 import zarr
+from anndata.experimental import write_elem
 from cellpose import models
 from cellpose.core import use_gpu
 
 import fractal_tasks_core
 from fractal_tasks_core.lib_pyramid_creation import build_pyramid
+from fractal_tasks_core.lib_regions_of_interest import (
+    array_to_bounding_box_table,
+)
 from fractal_tasks_core.lib_regions_of_interest import (
     convert_ROI_table_to_indices,
 )
+from fractal_tasks_core.lib_remove_FOV_overlaps import (
+    get_overlapping_pairs_3D,
+)
 from fractal_tasks_core.lib_zattrs_utils import extract_zyx_pixel_sizes
 from fractal_tasks_core.lib_zattrs_utils import rescale_datasets
 
@@ -121,6 +129,7 @@ def cellpose_segmentation(
     flow_threshold: float = 0.4,
     model_type: str = "nuclei",
     ROI_table_name: str = "FOV_ROI_table",
+    bounding_box_ROI_table_name: str = None,
 ) -> Dict[str, Any]:
     """
     Example inputs:
@@ -136,8 +145,8 @@ def cellpose_segmentation(
     # Set input path
     if len(input_paths) > 1:
         raise NotImplementedError
-    in_path = input_paths[0]
-    zarrurl = (in_path.parent.resolve() / component).as_posix() + "/"
+    in_path = input_paths[0].parent
+    zarrurl = (in_path.resolve() / component).as_posix() + "/"
     logger.info(zarrurl)
 
     # Read useful parameters from metadata
@@ -166,6 +175,9 @@ def cellpose_segmentation(
         f"{zarrurl}.zattrs", level=0
     )
 
+    actual_res_pxl_sizes_zyx = extract_zyx_pixel_sizes(
+        f"{zarrurl}.zattrs", level=labeling_level
+    )
     # Create list of indices for 3D FOVs spanning the entire Z direction
     list_indices = convert_ROI_table_to_indices(
         ROI_table,
@@ -308,6 +320,10 @@ def cellpose_segmentation(
 
     # Iterate over ROIs
     num_ROIs = len(list_indices)
+
+    if bounding_box_ROI_table_name:
+        bbox_dataframe_list = []
+
     logger.info(f"[{well_id}] Now starting loop over {num_ROIs} ROIs")
     for i_ROI, indices in enumerate(list_indices):
         # Define region
@@ -353,6 +369,20 @@ def cellpose_segmentation(
                     f"but dtype={label_dtype}"
                 )
 
+        if bounding_box_ROI_table_name:
+
+            bbox_df = array_to_bounding_box_table(
+                fov_mask, actual_res_pxl_sizes_zyx
+            )
+
+            bbox_dataframe_list.append(bbox_df)
+
+            overlap_list = []
+            for df in bbox_dataframe_list:
+                overlap_list.append(
+                    get_overlapping_pairs_3D(df, full_res_pxl_sizes_zyx)
+                )
+
         # Compute and store 0-th level to disk
         da.array(fov_mask).to_zarr(
             url=mask_zarr,
@@ -378,6 +408,23 @@ def cellpose_segmentation(
 
     logger.info(f"[{well_id}] End building pyramids, exit")
 
+    if bounding_box_ROI_table_name:
+        logger.info(f"[{well_id}] Writing bounding box table, exit")
+        # Concatenate all FOV dataframes
+        df_well = pd.concat(bbox_dataframe_list, axis=0, ignore_index=True)
+        df_well.index = df_well.index.astype(str)
+        # Convert all to float (warning: some would be int, in principle)
+        bbox_dtype = np.float32
+        df_well = df_well.astype(bbox_dtype)
+        # Convert to anndata
+        bbox_table = ad.AnnData(df_well, dtype=bbox_dtype)
+        # Write to zarr group
+        group_tables = zarr.group(f"{in_path}/{component}/tables/")
+        write_elem(group_tables, bounding_box_ROI_table_name, bbox_table)
+        logger.info(
+            f"[{in_path}/{component}/tables/{bounding_box_ROI_table_name}"
+        )
+
     return {}
 
 
@@ -402,6 +449,7 @@ class TaskArguments(BaseModel):
         flow_threshold: float = 0.4
         model_type: str = "nuclei"
         ROI_table_name: str = "FOV_ROI_table"
+        bounding_box_ROI_table_name: Optional[str] = None
 
     run_fractal_task(
         task_function=cellpose_segmentation, TaskArgsModel=TaskArguments

fractal_tasks_core/lib_regions_of_interest.py

@@ -285,3 +285,41 @@ def _inspect_ROI_table(
         print("Something went wrong in convert_ROI_table_to_indices\n", str(e))
 
     return df
+
+
+def array_to_bounding_box_table(
+    mask_array: np.ndarray, pxl_sizes_zyx: List[float]
+) -> pd.DataFrame:
+
+    """
+    Description
+
+    :param dummy: this is just a placeholder
+    :type dummy: int
+    """
+
+    labels = np.unique(mask_array)
+    labels = labels[labels > 0]
+    elem_list = []
+    for label in labels:
+        label_match = np.where(mask_array == label)
+        zmin, ymin, xmin = np.min(label_match, axis=1) * pxl_sizes_zyx
+        zmax, ymax, xmax = (np.max(label_match, axis=1) + 1) * pxl_sizes_zyx
+
+        length_x = xmax - xmin
+        length_y = ymax - ymin
+        length_z = zmax - zmin
+        elem_list.append((xmin, ymin, zmin, length_x, length_y, length_z))
+
+    df_columns = [
+        "x_micrometer",
+        "y_micrometer",
+        "z_micrometer",
+        "len_x_micrometer",
+        "len_y_micrometer",
+        "len_z_micrometer",
+    ]
+
+    ann_df = pd.DataFrame(np.array(elem_list), columns=df_columns)
+
+    return ann_df

fractal_tasks_core/lib_remove_FOV_overlaps.py

@@ -48,6 +48,27 @@ def is_overlapping_2D(box1, box2, tol=0):
     return overlap_x and overlap_y
 
 
+def is_overlapping_3D(box1, box2, tol=0):
+    """
+    Based on https://stackoverflow.com/a/70023212/19085332
+
+    box: (xmin, ymin, zmin, xmax, ymax, zmax)
+
+    :param dummy: this is just a placeholder
+    :type dummy: int
+    """
+    overlap_x = is_overlapping_1D(
+        [box1[0], box1[3]], [box2[0], box2[3]], tol=tol
+    )
+    overlap_y = is_overlapping_1D(
+        [box1[1], box1[4]], [box2[1], box2[4]], tol=tol
+    )
+    overlap_z = is_overlapping_1D(
+        [box1[2], box1[5]], [box2[2], box2[5]], tol=tol
+    )
+    return overlap_x and overlap_y and overlap_z
+
+
 def get_overlapping_pair(tmp_df, tol=0):
     """
     Description
@@ -66,6 +87,49 @@ def get_overlapping_pair(tmp_df, tol=0):
     return False
 
 
+def get_overlapping_pairs_3D(tmp_df, pixel_sizes):
+    """
+    Description
+
+    :param dummy: this is just a placeholder
+    :type dummy: int
+    """
+    # NOTE: here we use positional indices (i.e. starting from 0)
+    tol = 10e-10
+    if tol > min(pixel_sizes) / 1e3:
+        raise Exception(f"{tol=} but {pixel_sizes=}")
+    new_tmp_df = tmp_df.copy()
+
+    new_tmp_df["x_micrometer_max"] = (
+        new_tmp_df["x_micrometer"] + new_tmp_df["len_x_micrometer"]
+    )
+    new_tmp_df["y_micrometer_max"] = (
+        new_tmp_df["y_micrometer"] + new_tmp_df["len_y_micrometer"]
+    )
+    new_tmp_df["z_micrometer_max"] = (
+        new_tmp_df["z_micrometer"] + new_tmp_df["len_z_micrometer"]
+    )
+
+    new_tmp_df.drop(labels=["len_x_micrometer"], axis=1, inplace=True)
+    new_tmp_df.drop(labels=["len_y_micrometer"], axis=1, inplace=True)
+    new_tmp_df.drop(labels=["len_z_micrometer"], axis=1, inplace=True)
+    num_lines = len(new_tmp_df.index)
+    overlapping_list = []
+    # pos_ind_1 and pos_ind_2 are labels value
+    for pos_ind_1 in range(num_lines):
+        for pos_ind_2 in range(pos_ind_1):
+            if is_overlapping_3D(
+                new_tmp_df.iloc[pos_ind_1], new_tmp_df.iloc[pos_ind_2], tol=tol
+            ):
+                # we accumulate tuples of overlapping labels
+                overlapping_list.append((pos_ind_1, pos_ind_2))
+    if len(overlapping_list) > 0:
+        raise ValueError(
+            f"{overlapping_list} " f"List of pair of bounding box overlaps"
+        )
+    return overlapping_list
+
+
 def remove_FOV_overlaps(df):
     """
     Description