Add tracks exporting to zarr, napari, and trackmate. Some bug f… (

#75)

* Add tracks exporting to `zarr`, `napari`, and `trackmate`. Some bug fixes.

* Address review comments

* Refactor `export_tracks` and add tests for it

* Fix circular import

* add `recursive-include`

MANIFEST.in

@@ -1,3 +1,3 @@
 include README.md
 include LICENSE
-include ultrack/widgets/ultrackwidget/resources/*.json
+recursive-include ultrack/widgets/ultrackwidget/resources *.json

setup.cfg

@@ -44,6 +44,7 @@ install_requires =
     httpx >= 0.26.0
     websockets >= 12.0
     qtawesome >= 1.3.1
+    pydot >= 2.0.0
 
 [options.extras_require]
 testing =

ultrack/__init__.py

@@ -16,6 +16,7 @@
 
 from ultrack.config.config import MainConfig, load_config
 from ultrack.core.export.ctc import to_ctc
+from ultrack.core.export.exporter import export_tracks_by_extension
 from ultrack.core.export.trackmate import to_trackmate
 from ultrack.core.export.tracks_layer import to_tracks_layer
 from ultrack.core.export.zarr import tracks_to_zarr

ultrack/api/__init__.py

@@ -1,3 +1,2 @@
+from ultrack.api.database import Experiment, ExperimentStatus
 from ultrack.api.main import start_server
-from ultrack.api.database import Experiment
-from ultrack.api.database import ExperimentStatus

ultrack/api/database.py

@@ -7,7 +7,7 @@
 
 import sqlalchemy as sqla
 from pydantic import BaseModel, Json, validator
-from sqlalchemy import JSON, Column, DateTime, Enum, Integer, String, Text
+from sqlalchemy import JSON, Column, Enum, Integer, String, Text
 from sqlalchemy.orm import declarative_base, sessionmaker
 
 from ultrack import MainConfig
@@ -262,8 +262,6 @@ def update_experiment(experiment: Experiment) -> None:
     session.close()
 
 
-
-
 def get_experiment(id: int) -> Experiment:
     """Get an experiment from the database.
 

ultrack/api/main.py

@@ -12,14 +12,16 @@
 def _in_notebook():
     try:
         from IPython import get_ipython
-        if 'IPKernelApp' not in get_ipython().config:  # pragma: no cover
+
+        if "IPKernelApp" not in get_ipython().config:  # pragma: no cover
             return False
     except ImportError:
         return False
     except AttributeError:
         return False
     return True
 
+
 def start_server(
     api_results_path: Union[Path, str, None] = None,
     ultrack_data_config: Union[MainConfig, None] = None,
@@ -40,8 +42,10 @@ def start_server(
         os.environ["ULTRACK_DATA_CONFIG"] = ultrack_data_config.json()
 
     if _in_notebook():
+
         def start_in_notebook():
             uvicorn.run(app.app, host=host, port=port)
+
         Process(target=start_in_notebook).start()
     else:
         uvicorn.run(app.app, host=host, port=port)

ultrack/core/export/_test/test_exporter.py

@@ -0,0 +1,42 @@
+from pathlib import Path
+
+from ultrack import MainConfig, export_tracks_by_extension
+
+
+def test_exporter(tracked_database_mock_data: MainConfig, tmp_path: Path) -> None:
+    file_ext_list = [".xml", ".csv", ".zarr", ".dot", ".json"]
+    last_modified_time = {}
+    for file_ext in file_ext_list:
+        tmp_file = tmp_path / f"tracks{file_ext}"
+        export_tracks_by_extension(tracked_database_mock_data, tmp_file)
+
+        # assert file exists
+        assert (tmp_path / f"tracks{file_ext}").exists()
+        # assert file size is not zero
+        assert (tmp_path / f"tracks{file_ext}").stat().st_size > 0
+
+        # store last modified time
+        last_modified_time[str(tmp_file)] = tmp_file.stat().st_mtime
+
+    # loop again testing overwrite=False
+    for file_ext in file_ext_list:
+        tmp_file = tmp_path / f"tracks{file_ext}"
+        try:
+            export_tracks_by_extension(
+                tracked_database_mock_data, tmp_file, overwrite=False
+            )
+            assert False, "FileExistsError should be raised"
+        except FileExistsError:
+            pass
+
+    # loop again testing overwrite=True
+    for file_ext in file_ext_list:
+        tmp_file = tmp_path / f"tracks{file_ext}"
+        export_tracks_by_extension(tracked_database_mock_data, tmp_file, overwrite=True)
+
+        # assert file exists
+        assert (tmp_path / f"tracks{file_ext}").exists()
+        # assert file size is not zero
+        assert (tmp_path / f"tracks{file_ext}").stat().st_size > 0
+
+        assert last_modified_time[str(tmp_file)] != tmp_file.stat().st_mtime

ultrack/core/export/exporter.py

@@ -0,0 +1,74 @@
+import json
+from pathlib import Path
+from typing import Union
+
+import networkx as nx
+
+from ultrack.config import MainConfig
+from ultrack.core.export import (
+    to_networkx,
+    to_trackmate,
+    to_tracks_layer,
+    tracks_to_zarr,
+)
+
+
+def export_tracks_by_extension(
+    config: MainConfig, filename: Union[str, Path], overwrite: bool = False
+) -> None:
+    """
+    Export tracks to a file given the file extension.
+
+    Supported file extensions are .xml, .csv, .zarr, .dot, and .json.
+    - `.xml` exports to a TrackMate compatible XML file.
+    - `.csv` exports to a CSV file.
+    - `.zarr` exports the tracks to dense segments in a `zarr` array format.
+    - `.dot` exports to a Graphviz DOT file.
+    - `.json` exports to a networkx JSON file.
+
+    Parameters
+    ----------
+    filename : str or Path
+        The name of the file to save the tracks to.
+    config : MainConfig
+        The configuration object.
+    overwrite : bool, optional
+        Whether to overwrite the file if it already exists, by default False.
+
+    See Also
+    --------
+    to_trackmate :
+        Export tracks to a TrackMate compatible XML file.
+    to_tracks_layer :
+        Export tracks to a CSV file.
+    tracks_to_zarr :
+        Export tracks to a `zarr` array.
+    to_networkx :
+        Export tracks to a networkx graph.
+    """
+    if Path(filename).exists() and not overwrite:
+        raise FileExistsError(
+            f"File {filename} already exists. Set `overwrite=True` to overwrite the file"
+        )
+
+    file_ext = Path(filename).suffix
+    if file_ext.lower() == ".xml":
+        to_trackmate(config, filename, overwrite=True)
+    elif file_ext.lower() == ".csv":
+        df, _ = to_tracks_layer(config, include_parents=True)
+        df.to_csv(filename, index=False)
+    elif file_ext.lower() == ".zarr":
+        df, _ = to_tracks_layer(config)
+        tracks_to_zarr(config, df, filename, overwrite=True)
+    elif file_ext.lower() == ".dot":
+        G = to_networkx(config)
+        nx.drawing.nx_pydot.write_dot(G, filename)
+    elif file_ext.lower() == ".json":
+        G = to_networkx(config)
+        json_data = nx.node_link_data(G)
+        with open(filename, "w") as f:
+            json.dump(json_data, f)
+    else:
+        raise ValueError(
+            f"Unknown file extension: {file_ext}. Supported extensions are .xml, .csv, .zarr, .dot, and .json."
+        )

ultrack/core/export/networkx.py

@@ -91,5 +91,5 @@ def to_networkx(
     nx.DiGraph
         Networkx graph.
     """
-    df = to_tracks_layer(config)
+    df, _ = to_tracks_layer(config)
     return tracks_layer_to_networkx(df, children_to_parent)

ultrack/core/export/trackmate.py

@@ -83,7 +83,7 @@ def tracks_layer_to_trackmate(
             <ImageData filename="None" folder="None" width="0" height="0" depth="0" nslices="1" nframes="2" pixelwidth="1.0" pixelheight="1.0" voxeldepth="1.0" timeinterval="1.0"/>
         </Settings>
     </TrackMate>
-    """
+    """  # noqa: E501
     tracks_df["id"] = tracks_df["id"].astype(int)
     if not tracks_df["id"].is_unique:
         raise ValueError("The 'id' column must be unique.")

ultrack/core/tracker.py

@@ -7,6 +7,7 @@
 import zarr
 from numpy.typing import ArrayLike
 
+from ultrack import export_tracks_by_extension
 from ultrack.config import MainConfig
 from ultrack.core.export import (
     to_ctc,
@@ -147,3 +148,8 @@ def to_napari(self, *args, **kwargs) -> Tuple[pd.DataFrame, Dict]:
         self._assert_solved()
         tracks_df, graph = to_tracks_layer(self.config, *args, **kwargs)
         return tracks_df, graph
+
+    @functools.wraps(to_tracks_layer)
+    def export_by_extension(self, filename: str, overwrite: bool = False) -> None:
+        self._assert_solved()
+        export_tracks_by_extension(self.config, filename, overwrite=overwrite)

ultrack/widgets/ultrackwidget/ultrackwidget.py

@@ -18,6 +18,7 @@
     QGroupBox,
     QHBoxLayout,
     QLabel,
+    QMessageBox,
     QPushButton,
     QScrollArea,
     QSizePolicy,
@@ -27,14 +28,18 @@
     QWidget,
 )
 
-from ultrack import MainConfig
+from ultrack import MainConfig, export_tracks_by_extension
 from ultrack.widgets.ultrackwidget.components.button_workflow_config import (
     ButtonWorkflowConfig,
 )
 from ultrack.widgets.ultrackwidget.components.emitting_stream import EmittingStream
 from ultrack.widgets.ultrackwidget.data_forms import DataForms
 from ultrack.widgets.ultrackwidget.utils import UltrackInput
-from ultrack.widgets.ultrackwidget.workflows import UltrackWorkflow, WorkflowChoice
+from ultrack.widgets.ultrackwidget.workflows import (
+    UltrackWorkflow,
+    WorkflowChoice,
+    WorkflowStage,
+)
 
 LOG = logging.getLogger(__name__)
 
@@ -106,9 +111,54 @@ def _init_ui(self) -> None:
         self._add_run_button(layout)
         self._add_output_area(layout)
         self._add_cancel_button(layout)
+        self._add_bt_export_tracks(layout)
 
         layout.addStretch()
 
+    def _add_bt_export_tracks(self, layout: QVBoxLayout) -> None:
+        """
+        Add the export tracks button to the layout.
+
+        Parameters
+        ----------
+        layout : QVBoxLayout
+            The layout to which the export tracks button will be added.
+        """
+        self._add_spacer(layout, 10)
+        self._bt_export = QPushButton("Export tracks")
+        self._bt_export.setEnabled(False)
+        layout.addWidget(self._bt_export)
+
+    def export_tracks(self):
+        file_dialog = QFileDialog()
+        file_dialog.setFileMode(QFileDialog.AnyFile)
+        file_dialog.setAcceptMode(QFileDialog.AcceptSave)
+        file_dialog.setNameFilter(
+            "Napari Tracks (*.csv);;"
+            "Trackmate (*.xml);;"
+            "Zarr segments (*.zarr);;"
+            "NetworkX (*.dot);;"
+            "NetworkX (*.json)"
+        )
+
+        if file_dialog.exec_():
+            file_name = file_dialog.selectedFiles()[0]
+            ext = file_dialog.selectedNameFilter().split("*.")[-1][:-1]
+
+            # add the extension if not present
+            if not file_name.endswith(ext):
+                file_name += f".{ext}"
+
+            config = self._data_forms.get_config()
+            export_tracks_by_extension(config, file_name, overwrite=True)
+
+            QMessageBox.information(
+                self,
+                "Export tracks",
+                f"Tracks exported to {file_name}",
+                QMessageBox.Ok,
+            )
+
     def _add_title(self, layout: QVBoxLayout) -> None:
         """
         Add the title and subtitle to the layout.
@@ -346,6 +396,7 @@ def _setup_signals(self) -> None:
         self._bt_toggle_settings.clicked.connect(self._on_toggle_settings)
         self._bt_save_settings.clicked.connect(self._on_save_settings)
         self._bt_load_settings.clicked.connect(self._on_load_settings)
+        self._bt_export.clicked.connect(self.export_tracks)
         self._bt_run.clicked.connect(self._on_run)
         self._bt_cancel.clicked.connect(self._cancel)
         self._cb_workflow.currentIndexChanged.connect(self._on_workflow_changed)
@@ -396,7 +447,9 @@ def _on_workflow_intermediate_results(self, layer: Layer):
 
     def _on_run_started(self):
         """Handle the start of the run worker."""
+        QApplication.setOverrideCursor(QCursor(Qt.WaitCursor))
         self._bt_run.setEnabled(False)
+        self._bt_export.setEnabled(False)
         self._bt_run.setText("Running...")
         self._bt_run.repaint()
         self.ctx_stdout_switcher.__enter__()
@@ -410,13 +463,17 @@ def _on_run_finished(self):
         """Handle the finish of the run worker."""
         self._bt_run.setEnabled(True)
         self._bt_run.setText("Run")
+        self._bt_export.setEnabled(
+            self.workflow.last_reached_stage == WorkflowStage.DONE
+        )
         self._bt_run.repaint()
         self.ctx_stdout_switcher.__exit__(None, None, None)
         self.ctx_stderr_switcher.__exit__(None, None, None)
         self._output.hide()
         self._bt_cancel.hide()
         self.main_group.setEnabled(True)
         self._current_worker = None
+        QApplication.restoreOverrideCursor()
 
     @thread_worker
     def _make_run_worker(
@@ -434,11 +491,9 @@ def _make_run_worker(
         Worker
             The worker to run the selected workflow.
         """
-        QApplication.setOverrideCursor(QCursor(Qt.WaitCursor))
         yield from self.workflow.run(
             config, workflow_choice, inputs, additional_options
         )
-        QApplication.restoreOverrideCursor()
 
     def _on_save_settings(self) -> None:
         """Handle the save settings button click event."""
@@ -447,7 +502,7 @@ def _on_save_settings(self) -> None:
             None,
             "Save TOML",
             "ultrack_settings.toml",
-            "JSON Files (*.toml);;All Files (*)",
+            "TOML Files (*.toml);;All Files (*)",
             options=options,
         )
         if file_name:

ultrack/widgets/ultrackwidget/workflows.py

@@ -174,6 +174,7 @@ def run(
         self.inputs_used = inputs
         self.config = config
         self.additional_options = additional_options
+        self.last_reached_stage = WorkflowStage.PREPROCESSING
         try:
             if stage == WorkflowStage.PREPROCESSING:
                 if workflow_choice != WorkflowChoice.MANUAL: