Updating ground-truth matching and adding default classifier (#202)

* added save_config function

* added volume attribute to node

* close_tracks_gaps now works for segments being a dask_array

* removed node._volume method, since precomputed area already exists

* added UltrackArray class to utils/array.py

* ran pre-commit

* added path to database as optional input parameter. If provided, the database is taken from the provided path. If not provided, the database is loaded from the path stored in config

* hierarchy widget, work in progress

* added non-linear slider for uniform addition of volumes'

* added documentation to ultrack-array and HierarchyVizWidget

* Update ultrack/widgets/ultrackwidget/ultrackwidget.py

Co-authored-by: Jordão Bragantini <[email protected]>

* Update ultrack/widgets/ultrackwidget/ultrackwidget.py

Co-authored-by: Jordão Bragantini <[email protected]>

* Update ultrack/utils/array.py

Co-authored-by: Jordão Bragantini <[email protected]>

* implementing Jordaos revisions

* revising PR

* renaming functions in ultrack-array

* reverting dirty history changes

* redid the changes lost in merge

* added indexing test for ultrack-array

* added test for hierarchy widget

* added option to predict node probabilities

* fixed spelling 'persistense -> persistence'

* adding match ground-truth CLI

* fixed ground-truth match test

* bypassing edge case where there isn't any competing segmentation

* replacing xgboost with catboost

* added additional shape sanity check

* fixing right subset for fit

* Update ultrack/cli/_test/test_cli.py

* added batch index option to gt matcher

* adding option to use ground for pinning ILP variables

* added kwargs option to main function

* Update ultrack/ml/classification.py

* Update ultrack/utils/array.py

Co-authored-by: Jordão Bragantini <[email protected]>

* move UltrackArray to separate file to prevent circular inputs

* implemented jordaos review comments

* precommit

* WIP adding new link features classifier

* added ml to pixi

* adding link classifications modules

* add additional checking

* adding link classifier debugging code

* remove overlap bug fixes

* updating docs and typing

* splitting testing

* fixing widget testing

* adding constant to layer name

* Removing constant 5

* refactoring ultrack array to dynamics load any column

* added testing to variable column attribute

* updated UI and added offset

* minor test and combobox changes

* renaming vol to num_pix

* add scale from ultrack metadata

* changing link df update probability query

* added missing default value to empty match gt alignment

* adding caching

* bug fix

* fix sqlalchemy query binding

* adding ultrack array gtlink column support

* fixed multi time point loading

* changed matching policy to use intersection

* adding template matching

* fixing ILP

* ILP with template matching

* ignoring catboost files

* fixing deprecations, logging and typing

---------

Co-authored-by: TeunHuijben <[email protected]>
Co-authored-by: Teun Huijben <[email protected]>
Co-authored-by: TeunHuijben <[email protected]>

Author: JoOkuma

Diff for: .gitignore

@@ -181,3 +181,6 @@ data.db
 # pixi environments
 .pixi
 *.egg-info
+
+# catboost files
+catboost_info

Diff for: pyproject.toml

@@ -72,6 +72,10 @@ test = [
     "napari[testing] >0.4.18",
     "pyqt5 >=5.15.4",
 ]
+ml = [
+    "catboost >=1.2.7,<1.3",
+    "scikit-learn >=1.6.0,<1.7",
+]
 
 [project.scripts]
 ultrack = "ultrack.cli.main:main"
@@ -120,6 +124,8 @@ uvicorn = ">=0.27.0.post1"
 websocket = ">=0.2.1"
 websockets = ">=12.0"
 zarr = ">=2.15.0,<3.0.0"
+scikit-learn = ">=1.6.0,<1.7"
+catboost = ">=1.2.7,<1.3"
 pyarrow = ">=16.1.0,<20"
 
 [tool.pixi.feature.cuda]
@@ -142,6 +148,10 @@ pytest-qt = ">=4.4.0,<5"
 pyqt = ">=5.15.9,<6"
 pytest-cov = ">=6.0.0,<7"
 
+[tool.pixi.feature.ml.dependencies]
+catboost = ">=1.2.7,<1.3"
+scikit-learn = ">=1.6.0,<1.7"
+
 [tool.pytest.ini_options]
 filterwarnings = [
     "ignore::DeprecationWarning:pkg_resources.*:",
@@ -154,7 +164,8 @@ ultrack = { path = ".", editable = true }
 default = { solve-group = "default" }
 cuda = { features = ["cuda"] }
 # docs = { features = ["docs"]}  # Current dependencies aren't compatible with pixi
-test = { features = ["test"], solve-group = "default" }
+ml = { features = ["ml"], solve-group = "default" }
+test = { features = ["test", "ml"], solve-group = "default" }
 
 [tool.pixi.feature.test.tasks]
 test = "pytest -v --color=yes --cov=ultrack --cov-report=html --durations=15 ."

Diff for: ultrack/cli/_test/test_cli.py

@@ -65,6 +65,12 @@ def test_segment(
                 "foreground",
                 "-cl",
                 "contours",
+                "-il",
+                "labels",
+                "-il",
+                "contours",
+                "-p",
+                "intensity_mean",
             ]
             + zarr_dataset_paths
         )
@@ -86,6 +92,50 @@ def test_link_with_images(
             ["link", "-cfg", str(instance_config_path), "-ow"] + zarr_dataset_paths[:2]
         )
 
+    def test_fit_and_add_probs(
+        self, instance_config_path: str, tmp_path: Path, zarr_dataset_paths: List[str]
+    ) -> None:
+        # required by match gt with model output
+        pytest.importorskip("catboost")
+        pytest.importorskip("sklearn")
+
+        model_path = tmp_path / "model.pkl"
+        new_cfg_path = tmp_path / "new_config.toml"
+
+        _run_command(
+            [
+                "match_gt",
+                "-cfg",
+                instance_config_path,
+                "-gl",
+                "labels",
+                "-om",
+                str(model_path),
+                "-oc",
+                str(new_cfg_path),
+                "--is-tracking",
+                "--is-segmentation",
+                "--persistence",
+            ]
+            + zarr_dataset_paths
+        )
+
+        # testing loading new config
+        load_config(new_cfg_path)
+
+        for var in ["nodes", "links"]:
+            _run_command(
+                [
+                    "add_probs",
+                    str(model_path),
+                    "-cfg",
+                    instance_config_path,
+                    "--persistence",
+                    "--var",
+                    var,
+                ]
+            )
+
     def test_solve(self, instance_config_path: str) -> None:
         with pytest.warns(UserWarning):
             # batch index with overwrite should trigger warning
@@ -155,7 +205,7 @@ def test_zarr_napari_export(
             ]
         )
 
-    @pytest.mark.parametrize("mode", ["solutions", "links", "all"])
+    @pytest.mark.parametrize("mode", ["gt", "solutions", "links", "all"])
     def test_clear_database(self, instance_config_path: str, mode: str) -> None:
         _run_command(
             [

Diff for: ultrack/cli/clear_database.py

@@ -4,11 +4,12 @@
 from ultrack.config.config import MainConfig
 from ultrack.core.database import clear_all_data
 from ultrack.core.linking.utils import clear_linking_data
+from ultrack.core.match_gt import clear_ground_truths
 from ultrack.core.solve.sqltracking import SQLTracking
 
 
 @click.command("clear_database")
-@click.argument("mode", type=click.Choice(["all", "links", "solutions"]))
+@click.argument("mode", type=click.Choice(["all", "links", "solutions", "gt"]))
 @config_option()
 def clear_database_cli(mode: str, config: MainConfig) -> None:
     """Cleans database content."""
@@ -20,5 +21,7 @@ def clear_database_cli(mode: str, config: MainConfig) -> None:
         clear_linking_data(database_path)
     elif mode == "solutions":
         SQLTracking.clear_solution_from_database(database_path)
+    elif mode == "gt":
+        clear_ground_truths(database_path)
     else:
         raise NotImplementedError(f"Clear database mode {mode} not implemented.")

Diff for: ultrack/cli/main.py

@@ -9,6 +9,8 @@
 from ultrack.cli.flow import add_flow_cli
 from ultrack.cli.labels_to_edges import labels_to_contours_cli
 from ultrack.cli.link import link_cli
+from ultrack.cli.match_gt import match_gt_cli
+from ultrack.cli.predict import add_probs_cli
 from ultrack.cli.segment import segmentation_cli
 from ultrack.cli.server import server_cli
 from ultrack.cli.solve import solve_cli
@@ -29,6 +31,8 @@ def main():
 main.add_command(export_cli)
 main.add_command(labels_to_contours_cli)
 main.add_command(link_cli)
+main.add_command(match_gt_cli)
+main.add_command(add_probs_cli)
 main.add_command(segmentation_cli)
 main.add_command(solve_cli)
 main.add_command(server_cli)

Diff for: ultrack/cli/match_gt.py

@@ -0,0 +1,176 @@
+import logging
+from pathlib import Path
+from typing import Optional, Sequence
+
+import click
+import cloudpickle
+import toml
+from napari.plugins import _initialize_plugins
+from napari.viewer import ViewerModel
+from rich.logging import RichHandler
+
+from ultrack.cli.segment import _get_layer_data
+from ultrack.cli.utils import (
+    batch_index_option,
+    config_option,
+    napari_reader_option,
+    overwrite_option,
+    paths_argument,
+    persistence_option,
+)
+from ultrack.config import MainConfig
+from ultrack.core.match_gt import match_to_ground_truth
+from ultrack.ml.classification import fit_links_prob, fit_nodes_prob
+
+LOG = logging.getLogger(__name__)
+LOG.setLevel(logging.INFO)
+LOG.addHandler(RichHandler())
+
+
+@click.command("match_gt")
+@paths_argument()
+@napari_reader_option()
+@config_option()
+@click.option(
+    "--ground-truth-layer",
+    "-gl",
+    required=False,
+    type=str,
+    default=None,
+    help="Ground-truth layer index on napari.",
+)
+@click.option(
+    "--output-model",
+    "-om",
+    type=click.Path(dir_okay=False, path_type=Path),
+    required=False,
+    default=None,
+    help="Optional output model file path.",
+)
+@click.option(
+    "--output-config",
+    "-oc",
+    type=click.Path(dir_okay=False, path_type=Path),
+    help="Optional output config file path.",
+)
+@click.option(
+    "--is-segmentation",
+    is_flag=True,
+    type=bool,
+    default=False,
+    help="Indicates ground-truth are fully curated segmentation masks. "
+    "When activated different costs are used for insertions and deletions.",
+)
+@click.option(
+    "--is-tracking",
+    is_flag=True,
+    type=bool,
+    default=False,
+    help="Indicates ground-truth are tracking instances results.",
+)
+@click.option(
+    "--is-dense",
+    is_flag=True,
+    type=bool,
+    default=False,
+    help="Indicates ground-truth are dense annotations (everything is annotated).",
+)
+@click.option(
+    "--insert-prob",
+    is_flag=True,
+    type=bool,
+    default=False,
+    help="Insert estimated probabilities into the database.",
+)
+@batch_index_option()
+@overwrite_option()
+@persistence_option()
+def match_gt_cli(
+    paths: Sequence[Path],
+    reader_plugin: str,
+    config: MainConfig,
+    ground_truth_layer: Optional[str],
+    output_model: Optional[Path],
+    output_config: Optional[Path],
+    is_segmentation: bool,
+    is_tracking: bool,
+    is_dense: bool,
+    insert_prob: bool,
+    batch_index: Optional[int],
+    overwrite: bool,
+    persistence: bool,
+) -> None:
+    """
+    Match ground-truth labels to the segmentation/tracking database.
+    """
+
+    if output_model is not None and output_model.exists() and not overwrite:
+        raise FileExistsError(
+            f"Output model {output_model} already exists. Use --overwrite to overwrite."
+        )
+
+    if output_config is not None:
+        if not is_segmentation:
+            raise ValueError(
+                "Output config is only available for segmentation ground-truth `--is-segmentation`."
+            )
+
+        if output_config.exists() and not overwrite:
+            raise FileExistsError(
+                f"Output config {output_config} already exists. Use --overwrite to overwrite."
+            )
+
+    # Data loading
+    _initialize_plugins()
+
+    viewer = ViewerModel()
+    viewer.open(path=paths, plugin=reader_plugin)
+
+    if ground_truth_layer is None:
+        if len(viewer.layers) > 1:
+            raise ValueError(
+                "Multiple layers found, please specify `--ground-truth-layer`."
+            )
+        else:
+            ground_truth_layer = viewer.layers[0].name
+
+    gt = _get_layer_data(viewer, ground_truth_layer)
+
+    # Match ground-truth to database
+    gt_df, new_config = match_to_ground_truth(
+        config=config,
+        gt_labels=gt,
+        scale=config.data_config.metadata.get("scale"),
+        is_segmentation=is_segmentation,
+        optimize_config=True,
+        batch_index=batch_index,
+    )
+
+    if output_config is not None:
+        LOG.info("Estimated new config: %s", new_config)
+        LOG.info("Saving new config to %s", output_config)
+        with open(output_config, "w") as f:
+            toml.dump(new_config.model_dump(by_alias=True), f)
+
+    if insert_prob or output_model is not None:
+        model = fit_nodes_prob(
+            config,
+            gt_df["gt_track_id"],
+            persistence_features=persistence,
+            insert_prob=insert_prob,
+            remove_no_overlap=not is_dense,
+        )
+
+        if is_tracking:
+            link_model = fit_links_prob(
+                config,
+                gt_df["gt_track_id"],
+                persistence_features=persistence,
+                insert_prob=insert_prob,
+            )
+            model = {"nodes": model, "links": link_model}
+
+        if output_model is not None:
+            LOG.info("Saving model to %s", output_model)
+            with open(output_model, "wb") as f:
+                cloudpickle.dump(model, f)

Diff for: ultrack/cli/predict.py

@@ -0,0 +1,41 @@
+from pathlib import Path
+from typing import Literal
+
+import click
+from cloudpickle import load
+
+from ultrack.cli.utils import config_option, persistence_option
+from ultrack.config import MainConfig
+from ultrack.ml.classification import predict_links_prob, predict_nodes_prob
+
+
+@click.command("add_probs")
+@click.argument("classif_pickle_path", type=click.Path(exists=True, path_type=Path))
+@click.option(
+    "--var",
+    type=click.Choice(["nodes", "links", "divisions", "appearances", "disappearances"]),
+    default="nodes",
+    help="Variable to assign probabilities.",
+)
+@config_option()
+@persistence_option()
+def add_probs_cli(
+    classif_pickle_path: Path,
+    var: Literal["nodes", "links", "divisions", "appearances", "disappearances"],
+    config: MainConfig,
+    persistence: bool,
+) -> None:
+    """Predicts and adds nodes' probabilities to the database."""
+
+    with open(classif_pickle_path, "rb") as f:
+        classifier = load(f)
+        if isinstance(classifier, dict):
+            classifier = classifier[var]
+
+    if var == "nodes":
+        predict_nodes_prob(config, classifier, persistence_features=persistence)
+    elif var == "links":
+        predict_links_prob(config, classifier, persistence_features=persistence)
+    else:
+        # TODO add edges and other probabilities
+        raise NotImplementedError(f"Variable {var} not implemented.")