Re-factor component labeling for usage with CLI

Author: schilling40

flamingo_tools/segmentation/postprocessing.py

@@ -363,110 +363,43 @@ def graph_connected_components(coords: dict, max_edge_distance: float, min_compo
 
 def components_sgn(
     table: pd.DataFrame,
-    keyword: str = "distance_nn100",
-    threshold_erode: Optional[float] = None,
     min_component_length: int = 50,
     max_edge_distance: float = 30,
-    iterations_erode: int = 0,
-    postprocess_threshold: Optional[float] = None,
-    postprocess_components: Optional[List[int]] = None,
 ) -> List[List[int]]:
     """Eroding the SGN segmentation.
 
     Args:
         table: Dataframe of segmentation table.
-        keyword: Keyword of the dataframe column for erosion.
-        threshold_erode: Threshold of column value after erosion step with spatial statistics.
         min_component_length: Minimal length for filtering out connected components.
         max_edge_distance: Maximal distance in micrometer between points to create edges for connected components.
-        iterations_erode: Number of iterations for erosion.
-        postprocess_threshold: Post-process graph connected components by searching for points closer than threshold.
-        postprocess_components: Post-process specific graph connected components ([0] for largest component only).
 
     Returns:
         Subgraph components as lists of label_ids of dataframe.
     """
-    if keyword not in table:
-        distance_avg = nearest_neighbor_distance(table, n_neighbors=100)
-        table.loc[:, keyword] = list(distance_avg)
-
     centroids = list(zip(table["anchor_x"], table["anchor_y"], table["anchor_z"]))
     labels = [int(i) for i in list(table["label_id"])]
-
-    distance_nn = list(table[keyword])
-    distance_nn.sort()
-
-    if len(table) < 20000:
-        min_cells = None
-        average_dist = int(distance_nn[int(len(table) * 0.8)])
-        threshold = threshold_erode if threshold_erode is not None else average_dist
-    else:
-        min_cells = 20000
-        threshold = threshold_erode if threshold_erode is not None else 40
-
-    if iterations_erode != 0 and iterations_erode is not None:
-        print(f"Using threshold of {threshold} micrometer for eroding segmentation with keyword {keyword}.")
-        new_subset = erode_subset(table.copy(), iterations=iterations_erode,
-                                  threshold=threshold, min_cells=min_cells, keyword=keyword)
-    else:
-        new_subset = table.copy()
-
-    # create graph from coordinates of eroded subset
-    centroids_subset = list(zip(new_subset["anchor_x"], new_subset["anchor_y"], new_subset["anchor_z"]))
-    labels_subset = [int(i) for i in list(new_subset["label_id"])]
     coords = {}
-    for index, element in zip(labels_subset, centroids_subset):
+    for index, element in zip(labels, centroids):
         coords[index] = element
 
     components, _ = graph_connected_components(coords, max_edge_distance, min_component_length)
 
-    # add original coordinates closer to eroded component than threshold
-    if postprocess_threshold is not None:
-        if postprocess_components is None:
-            pp_components = components
-        else:
-            pp_components = [components[i] for i in postprocess_components]
-
-        add_coords = []
-        for label_id, centr in zip(labels, centroids):
-            if label_id not in labels_subset:
-                add_coord = []
-                for comp_index, component in enumerate(pp_components):
-                    for comp_label in component:
-                        dist = math.dist(centr, centroids[comp_label - 1])
-                        if dist <= postprocess_threshold:
-                            add_coord.append([comp_index, label_id])
-                            break
-                if len(add_coord) != 0:
-                    add_coords.append(add_coord)
-        if len(add_coords) != 0:
-            for c in add_coords:
-                components[c[0][0]].append(c[0][1])
-
     return components
 
 
 def label_components_sgn(
     table: pd.DataFrame,
     min_size: int = 1000,
-    threshold_erode: Optional[float] = None,
     min_component_length: int = 50,
     max_edge_distance: float = 30,
-    iterations_erode: int = 0,
-    postprocess_threshold: Optional[float] = None,
-    postprocess_components: Optional[List[int]] = None,
 ) -> List[int]:
     """Label SGN components using graph connected components.
 
     Args:
         table: Dataframe of segmentation table.
         min_size: Minimal number of pixels for filtering small instances.
-        threshold_erode: Threshold of column value after erosion step with spatial statistics.
         min_component_length: Minimal length for filtering out connected components.
         max_edge_distance: Maximal distance in micrometer between points to create edges for connected components.
-        iterations_erode: Number of iterations for erosion.
-        postprocess_threshold: Post-process graph connected components by searching for points closer than threshold.
-        postprocess_components: Post-process specific graph connected components ([0] for largest component only).
 
     Returns:
         List of component label for each point in dataframe. 0 - background, then in descending order of size
@@ -476,10 +409,8 @@ def label_components_sgn(
     entries_filtered = table[table.n_pixels < min_size]
     table = table[table.n_pixels >= min_size]
 
-    components = components_sgn(table, threshold_erode=threshold_erode, min_component_length=min_component_length,
-                                max_edge_distance=max_edge_distance, iterations_erode=iterations_erode,
-                                postprocess_threshold=postprocess_threshold,
-                                postprocess_components=postprocess_components)
+    components = components_sgn(table, min_component_length=min_component_length,
+                                max_edge_distance=max_edge_distance)
 
     # add size-filtered objects to have same initial length
     table = pd.concat([table, entries_filtered], ignore_index=True)

reproducibility/label_components/repro_label_components.py

@@ -1,7 +1,7 @@
 import argparse
 import json
 import os
-from typing import Optional
+from typing import List, Optional
 
 import pandas as pd
 from flamingo_tools.s3_utils import get_s3_path
@@ -67,86 +67,117 @@ def label_custom_components(tsv_table, custom_dict):
     return tsv_table
 
 
-def repro_label_components(
-    ddict: dict,
-    output_dir: str,
+def _load_json_as_list(ddict_path: str) -> List[dict]:
+    with open(ddict_path, "r") as f:
+        data = json.loads(f.read())
+    # ensure the result is always a list
+    return data if isinstance(data, list) else [data]
+
+
+def label_components_single(
+    table_path: str,
+    out_path: str,
+    cell_type: str = "sgn",
+    component_list: List[int] = [1],
+    max_edge_distance: float = 30,
+    min_component_length: int = 50,
+    min_size: int = 1000,
+    s3: bool = False,
     s3_credentials: Optional[str] = None,
     s3_bucket_name: Optional[str] = None,
     s3_service_endpoint: Optional[str] = None,
+    custom_dic: Optional[dict] = None,
+    **_
 ):
-    default_cell_type = "sgn"
-    default_component_list = [1]
-    default_iterations_erode = None
-    default_max_edge_distance = 30
-    default_min_length = 50
-    default_min_size = 1000
-    default_seg_channel = "SGN_v2"
-    default_threshold_erode = None
-
-    with open(ddict, "r") as myfile:
-        data = myfile.read()
-    param_dicts = json.loads(data)
-
-    for dic in param_dicts:
-        cochlea = dic["cochlea"]
-        print(f"\n{cochlea}")
-
-        cell_type = dic.get("cell_type", default_cell_type)
-        component_list = dic.get("component_list", default_component_list)
-        iterations_erode = dic.get("iterations_erode", default_iterations_erode)
-        max_edge_distance = dic.get("max_edge_distance", default_max_edge_distance)
-        min_component_length = dic.get("min_component_length", default_min_length)
-        min_size = dic.get("min_size", default_min_size)
-        table_name = dic.get("segmentation_channel", default_seg_channel)
-        threshold_erode = dic.get("threshold_erode", default_threshold_erode)
-
-        s3_path = os.path.join(f"{cochlea}", "tables", table_name, "default.tsv")
-        tsv_path, fs = get_s3_path(s3_path, bucket_name=s3_bucket_name,
+    """Process a single cochlea using one set of parameters or a custom_dic.
+    """
+    if s3:
+        tsv_path, fs = get_s3_path(table_path, bucket_name=s3_bucket_name,
                                    service_endpoint=s3_service_endpoint, credential_file=s3_credentials)
-        with fs.open(tsv_path, "r") as f:
-            table = pd.read_csv(f, sep="\t")
-
-        if "custom_dic" in list(dic.keys()):
-            print(len(table[table["component_labels"] == 1]))
-            tsv_table = label_custom_components(table, dic["custom_dic"])
+    with fs.open(tsv_path, "r") as f:
+        table = pd.read_csv(f, sep="\t")
+
+    if custom_dic is not None:
+        tsv_table = label_custom_components(table, custom_dic)
+    else:
+        if cell_type == "sgn":
+            tsv_table = label_components_sgn(table, min_size=min_size,
+                                             min_component_length=min_component_length,
+                                             max_edge_distance=max_edge_distance)
+        elif cell_type == "ihc":
+            tsv_table = label_components_ihc(table, min_size=min_size,
+                                             min_component_length=min_component_length,
+                                             max_edge_distance=max_edge_distance)
         else:
-            if cell_type == "sgn":
-                tsv_table = label_components_sgn(table, min_size=min_size,
-                                                 threshold_erode=threshold_erode,
-                                                 min_component_length=min_component_length,
-                                                 max_edge_distance=max_edge_distance,
-                                                 iterations_erode=iterations_erode)
-            elif cell_type == "ihc":
-                tsv_table = label_components_ihc(table, min_size=min_size,
-                                                 min_component_length=min_component_length,
-                                                 max_edge_distance=max_edge_distance)
-            else:
-                raise ValueError("Choose a supported cell type. Either 'sgn' or 'ihc'.")
+            raise ValueError("Choose a supported cell type. Either 'sgn' or 'ihc'.")
 
-        custom_comp = len(tsv_table[tsv_table["component_labels"].isin(component_list)])
-        print(f"Total {cell_type.upper()}s: {len(tsv_table)}")
-        if component_list == [1]:
-            print(f"Largest component has {custom_comp} {cell_type.upper()}s.")
-        else:
-            for comp in component_list:
-                print(f"Component {comp} has {len(tsv_table[tsv_table["component_labels"] == comp])} instances.")
-            print(f"Custom component(s) have {custom_comp} {cell_type.upper()}s.")
+    custom_comp = len(tsv_table[tsv_table["component_labels"].isin(component_list)])
+    print(f"Total {cell_type.upper()}s: {len(tsv_table)}")
+    if component_list == [1]:
+        print(f"Largest component has {custom_comp} {cell_type.upper()}s.")
+    else:
+        for comp in component_list:
+            num_instances = len(tsv_table[tsv_table["component_labels"] == comp])
+            print(f"Component {comp} has {num_instances} instances.")
+        print(f"Custom component(s) have {custom_comp} {cell_type.upper()}s.")
 
-        cochlea_str = "-".join(cochlea.split("_"))
-        table_str = "-".join(table_name.split("_"))
-        os.makedirs(output_dir, exist_ok=True)
-        out_path = os.path.join(output_dir, "_".join([cochlea_str, f"{table_str}.tsv"]))
+    tsv_table.to_csv(out_path, sep="\t", index=False)
 
-        tsv_table.to_csv(out_path, sep="\t", index=False)
+
+def repro_label_components(
+    output_path: str,
+    table_path: Optional[str] = None,
+    ddict: Optional[str] = None,
+    **kwargs
+):
+    """Wrapper function for labeling connected components using a segmentation table.
+    The function is used to distinguish between a passed parameter dictionary in JSON format
+    and the explicit setting of parameters.
+    """
+    if ddict is None:
+        label_components_single(table_path, output_path, **kwargs)
+    else:
+        param_dicts = _load_json_as_list(ddict)
+        for params in param_dicts:
+
+            cochlea = params["cochlea"]
+            print(f"\n{cochlea}")
+            seg_channel = params["segmentation_channel"]
+            table_path = os.path.join(f"{cochlea}", "tables", seg_channel, "default.tsv")
+
+            if os.path.isdir(output_path):
+                cochlea_str = "-".join(cochlea.split("_"))
+                table_str = "-".join(seg_channel.split("_"))
+                save_path = os.path.join(output_path, "_".join([cochlea_str, f"{table_str}.tsv"]))
+            else:
+                save_path = output_path
+            label_components_single(table_path=table_path, out_path=save_path, **params)
 
 
 def main():
     parser = argparse.ArgumentParser(
         description="Script to label segmentation using a segmentation table and graph connected components.")
 
-    parser.add_argument("-i", "--input", type=str, required=True, help="Input JSON dictionary.")
-    parser.add_argument("-o", "--output", type=str, required=True, help="Output directory.")
+    parser.add_argument("-o", "--output", type=str, required=True,
+                        help="Output path. Either directory or specific file.")
+
+    parser.add_argument("-i", "--input", type=str, default=None, help="Input path to segmentation table.")
+    parser.add_argument("-j", "--json", type=str, default=None, help="Input JSON dictionary.")
+
+    parser.add_argument("--cell_type", type=str, default="sgn",
+                        help="Cell type of segmentation. Either 'sgn' or 'ihc'.")
+
+    # options for post-processing
+    parser.add_argument("--min_size", type=int, default=1000,
+                        help="Minimal number of pixels for filtering small instances.")
+    parser.add_argument("--min_component_length", type=int, default=50,
+                        help="Minimal length for filtering out connected components.")
+    parser.add_argument("--max_edge_distance", type=float, default=30,
+                        help="Maximal distance in micrometer between points to create edges for connected components.")
+    parser.add_argument("-c", "--components", type=str, nargs="+", default=[1], help="List of connected components.")
 
+    # options for S3 bucket
+    parser.add_argument("--s3", action="store_true", help="Flag for using S3 bucket.")
     parser.add_argument("--s3_credentials", type=str, default=None,
                         help="Input file containing S3 credentials. "
                         "Optional if AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY were exported.")
@@ -158,8 +189,18 @@ def main():
     args = parser.parse_args()
 
     repro_label_components(
-        args.input, args.output,
-        args.s3_credentials, args.s3_bucket_name, args.s3_service_endpoint,
+        output_path=args.output,
+        table_path=args.input,
+        ddict=args.json,
+        cell_type=args.cell_type,
+        component_list=args.components,
+        max_edge_distance=args.max_edge_distance,
+        min_component_length=args.min_component_length,
+        min_size=args.min_size,
+        s3=args.s3,
+        s3_credentials=args.s3_credentials,
+        s3_bucket_name=args.s3_bucket_name,
+        s3_service_endpoint=args.s3_service_endpoint,
     )