145 option to export segmentation relabeled by track

pyproject.toml

@@ -31,7 +31,7 @@ classifiers = [
 
 dependencies =[
     "napari>=0.6.2,<1",
-    "funtracks>=1.7.0,<2",
+    "funtracks>=1.8.0-a2,<2",
     "appdirs>=1,<2",
     "numpy>=2,<3",
     "magicgui>=0.10.1",

src/motile_tracker/data_views/views_coordinator/groups.py

@@ -484,4 +484,5 @@ def _show_export_dialog(self, item: QListWidgetItem) -> None:
             tracks=self.tracks_viewer.tracks,
             name=group_name,
             nodes_to_keep=nodes_to_keep,
+            colormap=self.tracks_viewer.colormap,
         )

src/motile_tracker/data_views/views_coordinator/tracks_list.py

@@ -92,9 +92,12 @@ class TracksList(QGroupBox):
     """
 
     view_tracks = Signal(Tracks, str)
+    request_colormap = Signal()
 
     def __init__(self):
         super().__init__(title="Results List")
+
+        self.colormap = None
         self.file_dialog = QFileDialog()
         self.file_dialog.setFileMode(QFileDialog.Directory)
         self.file_dialog.setOption(QFileDialog.ShowDirsOnly, True)
@@ -178,11 +181,11 @@ def show_export_dialog(self, item: QListWidgetItem) -> None:
         widget: TracksButton = self.tracks_list.itemWidget(item)
         tracks: Tracks = widget.tracks
         name: str = widget.name.text()
+        self.request_colormap.emit()
+        colormap = self.colormap
 
         ExportDialog.show_export_dialog(
-            self,
-            tracks=tracks,
-            name=name,
+            self, tracks=tracks, name=name, colormap=colormap
         )
 
     def save_tracks(self, item: QListWidgetItem):

src/motile_tracker/data_views/views_coordinator/tracks_viewer.py

@@ -79,6 +79,7 @@ def __init__(
 
         self.tracks_list = TracksList()
         self.tracks_list.view_tracks.connect(self.update_tracks)
+        self.tracks_list.request_colormap.connect(self.set_colormap_to_trackslist)
         self.selected_track = None
         self.track_id_color = [0, 0, 0, 0]
         self.force = False
@@ -90,6 +91,10 @@ def __init__(
 
         self.viewer.dims.events.ndisplay.connect(self.update_selection)
 
+    def set_colormap_to_trackslist(self):
+        """Set the current colormap on the TracksList, so that it can be exported."""
+        self.tracks_list.colormap = self.colormap
+
     def set_keybinds(self):
         bind_keymap(self.viewer, KEYMAP, self)
 

src/motile_tracker/import_export/menus/export_dialog.py

@@ -1,21 +1,76 @@
 from pathlib import Path
 
+import napari
 from funtracks.data_model import Tracks
 from funtracks.import_export import export_to_csv
 from funtracks.import_export.export_to_geff import export_to_geff
 from qtpy.QtWidgets import (
+    QCheckBox,
+    QComboBox,
+    QDialog,
+    QDialogButtonBox,
     QFileDialog,
-    QInputDialog,
+    QLabel,
     QMessageBox,
+    QVBoxLayout,
 )
 
 
+class ExportTypeDialog(QDialog):
+    def __init__(self, parent=None, label: str = ""):
+        super().__init__(parent)
+        self.setWindowTitle("Select Export Type")
+
+        layout = QVBoxLayout(self)
+
+        if label:
+            layout.addWidget(QLabel(label))
+
+        self.export_type_combo = QComboBox()
+        self.export_type_combo.addItems(["CSV", "geff"])
+        layout.addWidget(self.export_type_combo)
+
+        self.relabel_checkbox = QCheckBox(
+            "Export segmentation relabeled by Tracklet ID"
+        )
+        layout.addWidget(self.relabel_checkbox)
+
+        buttons = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel)
+        layout.addWidget(buttons)
+
+        buttons.accepted.connect(self.accept)
+        buttons.rejected.connect(self.reject)
+
+        # Initial visibility
+        self._update_checkbox_visibility(self.export_type_combo.currentText())
+
+        # Update visibility when export type changes
+        self.export_type_combo.currentTextChanged.connect(
+            self._update_checkbox_visibility
+        )
+
+    def _update_checkbox_visibility(self, export_type: str):
+        self.relabel_checkbox.setVisible(export_type == "CSV")
+
+    @property
+    def export_type(self) -> str:
+        return self.export_type_combo.currentText()
+
+    @property
+    def relabel_by_tracklet_id(self) -> bool:
+        return self.relabel_checkbox.isChecked()
+
+
 class ExportDialog:
     """Handles exporting tracks to CSV or Geff."""
 
     @staticmethod
     def show_export_dialog(
-        parent, tracks: Tracks, name: str, nodes_to_keep: set[int] | None = None
+        parent,
+        tracks: Tracks,
+        name: str,
+        colormap: napari.utils.Colormap,
+        nodes_to_keep: set[int] | None = None,
     ):
         """
         Export tracks to CSV or Geff, with the option to export a subset of nodes only.
@@ -41,31 +96,66 @@ def show_export_dialog(
                 f"<p>Choose export format:</p>"
             )
 
-        export_type, ok = QInputDialog.getItem(
-            parent,
-            "Select Export Type",
-            label,
-            ["CSV", "geff"],
-            0,
-            False,
-        )
+        dialog = ExportTypeDialog(parent, label)
 
-        if not ok:
+        if dialog.exec_() != QDialog.Accepted:
             return False
 
+        export_type = dialog.export_type
+        relabel_by_tracklet_id = dialog.relabel_by_tracklet_id
+
         if export_type == "CSV":
-            file_dialog = QFileDialog(parent)
-            file_dialog.setFileMode(QFileDialog.AnyFile)
-            file_dialog.setAcceptMode(QFileDialog.AcceptSave)
-            file_dialog.setNameFilter("CSV files (*.csv)")
-            file_dialog.setDefaultSuffix("csv")
-            default_file = f"{name}_tracks.csv"
-            file_dialog.selectFile(str(Path.home() / default_file))
+            # CSV file dialog
+            csv_dialog = QFileDialog(parent, "Save to CSV")  # set title
+            csv_dialog.setFileMode(QFileDialog.AnyFile)
+            csv_dialog.setAcceptMode(QFileDialog.AcceptSave)
+            csv_dialog.setNameFilter("CSV files (*.csv)")
+            csv_dialog.setDefaultSuffix("csv")
+            default_csv_file = f"{name}_tracks.csv"
+            csv_dialog.selectFile(str(Path.home() / default_csv_file))
 
-            if file_dialog.exec_():
-                file_path = Path(file_dialog.selectedFiles()[0])
-                export_to_csv(tracks, file_path, nodes_to_keep, use_display_names=True)
-                return True
+            if not csv_dialog.exec_():
+                return False  # User canceled
+
+            file_path = Path(csv_dialog.selectedFiles()[0])
+            seg_path = None
+
+            # Optional segmentation dialog
+            if relabel_by_tracklet_id:
+                default_seg_path = file_path.with_suffix(".tif")
+                seg_dialog = QFileDialog(
+                    parent, "Save segmentation to TIF"
+                )  # set title
+                seg_dialog.setFileMode(QFileDialog.AnyFile)
+                seg_dialog.setAcceptMode(QFileDialog.AcceptSave)
+                seg_dialog.setNameFilter("TIF files (*.tif)")
+                seg_dialog.setDefaultSuffix("tif")
+                seg_dialog.selectFile(str(default_seg_path))
+
+                if not seg_dialog.exec_():
+                    return False  # User canceled
+
+                seg_path = Path(seg_dialog.selectedFiles()[0])
+
+            # Construct color_dict from colormap
+            nodes = list(tracks.graph.nodes())
+            track_ids = [tracks.get_track_id(node) for node in nodes]
+            colors = [colormap.map(tid) for tid in track_ids]
+            color_dict = {
+                **dict(zip(nodes, colors, strict=True)),
+                None: [0, 0, 0, 0],
+            }
+
+            export_to_csv(
+                tracks=tracks,
+                outfile=file_path,
+                color_dict=color_dict,
+                node_ids=nodes_to_keep,
+                use_display_names=True,
+                export_seg=relabel_by_tracklet_id,
+                seg_path=seg_path,
+            )
+            return True
 
         elif export_type == "geff":
             file_dialog = QFileDialog(parent, "Save as geff file")