Implement RF classification workflow

Author: constantinpape

flamingo_tools/classification/__init__.py

@@ -1 +1,2 @@
 from .classification_gui import run_classification_gui
+from .training_and_prediction import train_classifier, predict_classifier

flamingo_tools/classification/classification_gui.py

@@ -7,38 +7,29 @@
 import imageio.v3 as imageio
 import napari
 import numpy as np
-import pandas as pd
 
 from joblib import dump
 from magicgui import magic_factory
-from skimage.measure import regionprops_table
 
 import micro_sam.sam_annotator.object_classifier as classifier_util
 from micro_sam.object_classification import project_prediction_to_segmentation
 from micro_sam.sam_annotator._widgets import _generate_message
 
+from ..measurements import compute_object_measures_impl
+
 IMAGE_LAYER_NAME = None
 SEGMENTATION_LAYER_NAME = None
 FEATURES = None
 SEG_IDS = None
 CLASSIFIER = None
 LABELS = None
+FEATURE_SET = None
 
 
-# TODO refactor
 def _compute_features(segmentation, image):
-    features = pd.DataFrame(regionprops_table(
-        segmentation, image, properties=[
-            "label", "area", "axis_major_length", "axis_minor_length",
-            "equivalent_diameter_area", "euler_number", "extent",
-            "feret_diameter_max", "inertia_tensor_eigvals",
-            "intensity_max", "intensity_mean", "intensity_min",
-            "intensity_std", "moments_central",
-            "moments_weighted", "solidity",
-        ]
-    ))
-    seg_ids = features.label.values.astype(int)
-    features = features.drop(columns="label").values
+    features = compute_object_measures_impl(image, segmentation, feature_set=FEATURE_SET)
+    seg_ids = features.label_id.values.astype(int)
+    features = features.drop(columns="label_id").values
     return features, seg_ids
 
 
@@ -92,12 +83,29 @@ def _create_export_feature_widget(export_path: Optional[Path] = None) -> None:
     export_path = Path(export_path).with_suffix(".h5")
     with h5py.File(export_path, "a") as f:
         g = f.create_group(IMAGE_LAYER_NAME)
+        g.attrs["feature_set"] = FEATURE_SET
         g.create_dataset("features", data=features, compression="lzf")
         g.create_dataset("labels", data=labels, compression="lzf")
 
 
-def run_classification_gui(image_path, segmentation_path, image_name=None, segmentation_name=None):
-    global IMAGE_LAYER_NAME, SEGMENTATION_LAYER_NAME
+def run_classification_gui(
+    image_path: str,
+    segmentation_path: str,
+    image_name: Optional[str] = None,
+    segmentation_name: Optional[str] = None,
+    feature_set: str = "default",
+) -> None:
+    """Start the classification GUI.
+
+    Args:
+        image_path: The path to the image data.
+        segmentation_path: The path to the segmentation.
+        image_name: The name for the image layer. Will use the filename if not given.
+        segmentation_name: The name of the label layer with the segmentation.
+            Will use the filename if not given.
+        feature_set: The feature set to use. Refer to `flamingo_tools.measurements.FEATURE_FUNCTIONS` for details.
+    """
+    global IMAGE_LAYER_NAME, SEGMENTATION_LAYER_NAME, FEATURE_SET
 
     image = imageio.imread(image_path)
     segmentation = imageio.imread(segmentation_path)
@@ -107,6 +115,7 @@ def run_classification_gui(image_path, segmentation_path, image_name=None, segme
 
     IMAGE_LAYER_NAME = image_name
     SEGMENTATION_LAYER_NAME = segmentation_name
+    FEATURE_SET = feature_set
 
     viewer = napari.Viewer()
     viewer.add_image(image, name=image_name)

flamingo_tools/classification/training_and_prediction.py

@@ -1,10 +1,89 @@
+import multiprocessing as mp
+from typing import Optional, Sequence
 
+import h5py
+import numpy as np
+import pandas as pd
+from joblib import dump, load
+from sklearn.ensemble import RandomForestClassifier
 
-# TODO train a classifier on all features and labels stored in h5
-def train_classifier(feature_paths):
-    pass
+from ..measurements import compute_object_measures
 
 
-# TODO run prediction on a full cochlea
-def predict_classifier():
-    pass
+def train_classifier(feature_paths: Sequence[str], save_path: str, **rf_kwargs) -> None:
+    """Train a random forest classifier on features and labels that were exported via the classification GUI.
+
+    Args:
+        feature_paths: The path to the h5 files with features and labels.
+        save_path: Where to save the trained random forest.
+        rf_kwargs: Keyword arguments for creating the random forest.
+    """
+    features, labels = [], []
+    for path in feature_paths:
+        with h5py.File(path, "r") as f:
+            for name, group in f.items():
+                features.append(group["features"][:])
+                labels.append(group["labels"][:])
+
+    features = np.concatenate(features)
+    labels = np.concatenate(labels)
+
+    rf = RandomForestClassifier(**rf_kwargs)
+    rf.fit(features, labels)
+
+    dump(rf, save_path)
+
+
+def predict_classifier(
+    rf_path: str,
+    image_path: str,
+    segmentation_path: str,
+    feature_table_path: str,
+    segmentation_table_path: Optional[str],
+    image_key: Optional[str] = None,
+    segmentation_key: Optional[str] = None,
+    n_threads: Optional[int] = None,
+    feature_set: str = "default",
+) -> pd.DataFrame:
+    """Run prediction with a trained classifier on an input volume with associated segmentation.
+
+    Args:
+        rf_path: The path to the trained random forest.
+        image_path: The path to the image data.
+        segmentation_path: The path to the segmentation.
+        feature_table_path: The path for the features used for prediction.
+            The features will be computed and saved if this table does not exist.
+        segmentation_table_path: The path to the segmentation table (in MoBIE format).
+            It will be computed on the fly if it is not given.
+        image_key: The key / internal path for the image data. Not needed for tif data.
+        segmentation_key: The key / internal path for the segmentation data. Not needed for tif data.
+        n_threads: The number of threads for parallelization.
+        feature_set: The feature set to use. Refer to `flamingo_tools.measurements.FEATURE_FUNCTIONS` for details.
+
+    Returns:
+        A dataframe with the prediction. It contains the columns 'label_id', 'predictions' and
+            'probs-0', 'probs-1', ... . The latter columns contain the probabilities for the respective class.
+    """
+    compute_object_measures(
+        image_path=image_path,
+        segmentation_path=segmentation_path,
+        segmentation_table_path=segmentation_table_path,
+        output_table_path=feature_table_path,
+        image_key=image_key,
+        segmentation_key=segmentation_key,
+        n_threads=n_threads,
+        feature_set=feature_set,
+    )
+
+    features = pd.read_csv(feature_table_path, sep="\t")
+    label_ids = features.label_id.values
+    features = features.drop(columns=["label_id"]).values
+
+    rf = load(rf_path)
+    n_threads = mp.cpu_count() if n_threads is None else n_threads
+    rf.n_jobs_ = n_threads
+
+    probs = rf.predict_proba(features)
+    result = {"label_id": label_ids, "prediction": np.argmax(probs, axis=1)}
+    result.update({"probs-{i}": probs[:, i] for i in range(probs.shape[1])})
+    return pd.DataFrame(result)