From 182819f788ae5b05269f2d663839d1f1af485586 Mon Sep 17 00:00:00 2001 From: "Ilan F. S. Theodoro" Date: Wed, 3 Jul 2024 12:12:42 -0700 Subject: [PATCH] =?UTF-8?q?Add=20tracks=20exporting=20to=20`zarr`,=20`napa?= =?UTF-8?q?ri`,=20and=20`trackmate`.=20Some=20bug=20f=E2=80=A6=20(#75)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 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 | 2 +- setup.cfg | 1 + ultrack/__init__.py | 1 + ultrack/api/__init__.py | 3 +- ultrack/api/database.py | 4 +- ultrack/api/main.py | 6 +- ultrack/core/export/_test/test_exporter.py | 42 +++++++++++ ultrack/core/export/exporter.py | 74 +++++++++++++++++++ ultrack/core/export/networkx.py | 2 +- ultrack/core/export/trackmate.py | 2 +- ultrack/core/tracker.py | 6 ++ .../widgets/ultrackwidget/ultrackwidget.py | 65 ++++++++++++++-- ultrack/widgets/ultrackwidget/workflows.py | 1 + 13 files changed, 195 insertions(+), 14 deletions(-) create mode 100644 ultrack/core/export/_test/test_exporter.py create mode 100644 ultrack/core/export/exporter.py diff --git a/MANIFEST.in b/MANIFEST.in index 32efa63..1ebd88e 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,3 +1,3 @@ include README.md include LICENSE -include ultrack/widgets/ultrackwidget/resources/*.json +recursive-include ultrack/widgets/ultrackwidget/resources *.json \ No newline at end of file diff --git a/setup.cfg b/setup.cfg index e3db3cd..79eab98 100644 --- a/setup.cfg +++ b/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 = diff --git a/ultrack/__init__.py b/ultrack/__init__.py index 6484ad5..490bb7f 100644 --- a/ultrack/__init__.py +++ b/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 diff --git a/ultrack/api/__init__.py b/ultrack/api/__init__.py index 29d72d1..3fb5a15 100644 --- a/ultrack/api/__init__.py +++ b/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 \ No newline at end of file diff --git a/ultrack/api/database.py b/ultrack/api/database.py index 1959b34..28aa248 100644 --- a/ultrack/api/database.py +++ b/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. diff --git a/ultrack/api/main.py b/ultrack/api/main.py index 507cd04..fb44812 100644 --- a/ultrack/api/main.py +++ b/ultrack/api/main.py @@ -12,7 +12,8 @@ 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 @@ -20,6 +21,7 @@ def _in_notebook(): 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) diff --git a/ultrack/core/export/_test/test_exporter.py b/ultrack/core/export/_test/test_exporter.py new file mode 100644 index 0000000..8595b90 --- /dev/null +++ b/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 diff --git a/ultrack/core/export/exporter.py b/ultrack/core/export/exporter.py new file mode 100644 index 0000000..9583251 --- /dev/null +++ b/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." + ) diff --git a/ultrack/core/export/networkx.py b/ultrack/core/export/networkx.py index 9ee909e..66681cc 100644 --- a/ultrack/core/export/networkx.py +++ b/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) diff --git a/ultrack/core/export/trackmate.py b/ultrack/core/export/trackmate.py index 2e414ed..ba17acc 100644 --- a/ultrack/core/export/trackmate.py +++ b/ultrack/core/export/trackmate.py @@ -83,7 +83,7 @@ def tracks_layer_to_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.") diff --git a/ultrack/core/tracker.py b/ultrack/core/tracker.py index fa95f59..d3b0370 100644 --- a/ultrack/core/tracker.py +++ b/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) diff --git a/ultrack/widgets/ultrackwidget/ultrackwidget.py b/ultrack/widgets/ultrackwidget/ultrackwidget.py index 64f960c..247927b 100644 --- a/ultrack/widgets/ultrackwidget/ultrackwidget.py +++ b/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,6 +463,9 @@ 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) @@ -417,6 +473,7 @@ def _on_run_finished(self): 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: diff --git a/ultrack/widgets/ultrackwidget/workflows.py b/ultrack/widgets/ultrackwidget/workflows.py index 792972b..60c818b 100644 --- a/ultrack/widgets/ultrackwidget/workflows.py +++ b/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: