Continue track point button and other new features

src/motile_tracker/application_menus/editing_menu.py

@@ -1,7 +1,10 @@
 import napari
 from qtpy.QtWidgets import (
+    QButtonGroup,
     QGroupBox,
+    QHBoxLayout,
     QPushButton,
+    QRadioButton,
     QVBoxLayout,
     QWidget,
 )
@@ -15,8 +18,34 @@ def __init__(self, viewer: napari.Viewer):
 
         self.tracks_viewer = TracksViewer.get_instance(viewer)
         self.tracks_viewer.selected_nodes.list_updated.connect(self.update_buttons)
+        self.tracks_viewer.tracks_updated.connect(self._update_continue_tracks_box)
         layout = QVBoxLayout()
 
+        self.continue_tracks_box = QGroupBox("Point layer track mode")
+        self.continue_tracks_box.setMaximumHeight(60)
+        continue_tracks_layout = QHBoxLayout()
+        button_group = QButtonGroup()
+        self.continue_tracks_radio = QRadioButton("Continue tracks")
+        self.continue_tracks_radio.clicked.connect(
+            lambda: self.toggle_track_mode("continue")
+        )
+        self.new_track_radio = QRadioButton("Start new tracks")
+        self.new_track_radio.setChecked(True)
+        self.new_track_radio.clicked.connect(
+            lambda: self.toggle_track_mode("new track")
+        )
+        button_group.addButton(self.continue_tracks_radio)
+        button_group.addButton(self.new_track_radio)
+        continue_tracks_layout.addWidget(self.continue_tracks_radio)
+        continue_tracks_layout.addWidget(self.new_track_radio)
+
+        self.continue_tracks_box.setLayout(continue_tracks_layout)
+        if (
+            self.tracks_viewer.tracking_layers.seg_layer is not None
+            and self.tracks_viewer.tracking_layers.points_layer is not None
+        ):
+            self.continue_tracks_box.hide()
+
         node_box = QGroupBox("Edit Node(s)")
         node_box.setMaximumHeight(60)
         node_box_layout = QVBoxLayout()
@@ -63,13 +92,30 @@ def __init__(self, viewer: napari.Viewer):
         self.redo_btn = QPushButton("Redo (R)")
         self.redo_btn.clicked.connect(self.tracks_viewer.redo)
 
+        layout.addWidget(self.continue_tracks_box)
         layout.addWidget(node_box)
         layout.addWidget(edge_box)
         layout.addWidget(self.undo_btn)
         layout.addWidget(self.redo_btn)
 
         self.setLayout(layout)
-        self.setMaximumHeight(300)
+        self.setMaximumHeight(360)
+
+    def _update_continue_tracks_box(self):
+        """Show or hide the continue tracks box depending on the presence of the TracksLabels layer. If a TracksLabels layer is present, adding nodes with points is disabled and therefore the continue_tracks_button should be hidden."""
+
+        if (
+            self.tracks_viewer.tracking_layers.seg_layer is not None
+            and self.tracks_viewer.tracking_layers.points_layer is not None
+        ):
+            self.continue_tracks_box.hide()
+        else:
+            self.continue_tracks_box.show()
+
+    def toggle_track_mode(self, mode: str):
+        """Toggle the track mode (continue / new track on the tracks viewer)"""
+
+        self.tracks_viewer.track_mode = mode
 
     def update_buttons(self):
         """Set the buttons to enabled/disabled depending on the currently selected nodes"""

src/motile_tracker/application_menus/menu_widget.py

@@ -2,6 +2,7 @@
 from qtpy.QtWidgets import QScrollArea, QTabWidget, QVBoxLayout
 
 from motile_tracker.application_menus.editing_menu import EditingMenu
+from motile_tracker.data_views.views.view_3d import View3D
 from motile_tracker.data_views.views_coordinator.tracks_viewer import TracksViewer
 from motile_tracker.motile.menus.motile_widget import MotileWidget
 
@@ -16,9 +17,11 @@ def __init__(self, viewer: napari.Viewer):
 
         motile_widget = MotileWidget(viewer)
         editing_widget = EditingMenu(viewer)
+        view3D_widget = View3D(viewer)
 
         tabwidget = QTabWidget()
 
+        tabwidget.addTab(view3D_widget, "3D viewing")
         tabwidget.addTab(motile_widget, "Track with Motile")
         tabwidget.addTab(editing_widget, "Edit Tracks")
         tabwidget.addTab(tracks_viewer.tracks_list, "Results List")

src/motile_tracker/data_model/solution_tracks.py

@@ -11,6 +11,7 @@
     from collections.abc import Iterable
     from pathlib import Path
 
+    import napari
     import numpy as np
 
     from .tracks import Attrs, Node
@@ -74,6 +75,19 @@ def set_track_id(self, node: Node, value: int):
             self.track_id_to_node[value] = []
         self.track_id_to_node[value].append(node)
 
+    def get_lineage_id(self, node: Node) -> int:
+        """Return the track id value of the root node as lineage id"""
+
+        # go up the tree to identify the root node
+        root_node = node
+        while True:
+            predecessors = list(self.graph.predecessors(root_node))
+            if not predecessors:
+                break
+            root_node = predecessors[0]
+
+        return self.get_track_id(root_node)
+
     def _initialize_track_ids(self):
         self.max_track_id = 0
         self.track_id_to_node = {}
@@ -116,13 +130,30 @@ def _assign_tracklet_ids(self):
             track_id += 1
         self.max_track_id = track_id - 1
 
-    def export_tracks(self, outfile: Path | str):
+    def export_tracks(
+        self,
+        outfile: Path | str,
+        colormap: napari.utils.CyclicLabelColormap,
+    ):
         """Export the tracks from this run to a csv with the following columns:
-        t,[z],y,x,id,parent_id,track_id
+        t, [z], y, x, id, parent_id, track_id, lineage_id, color
         Cells without a parent_id will have an empty string for the parent_id.
         Whether or not to include z is inferred from self.ndim
+        Args:
+            outfile (Path | str): The path to the output file
+            colormap (napari.utils.CyclicLabelColormap): The colormap from which to infer the color by track id.
         """
-        header = ["t", "z", "y", "x", "id", "parent_id", "track_id"]
+        header = [
+            "t",
+            "z",
+            "y",
+            "x",
+            "id",
+            "parent_id",
+            "track_id",
+            "lineage_id",
+            "color",
+        ]
         if self.ndim == 3:
             header = [header[0]] + header[2:]  # remove z
         with open(outfile, "w") as f:
@@ -133,12 +164,16 @@ def export_tracks(self, outfile: Path | str):
                 track_id = self.get_track_id(node_id)
                 time = self.get_time(node_id)
                 position = self.get_position(node_id)
+                lineage_id = self.get_lineage_id(node_id)
+                color = colormap.map(track_id)[:3] * 255
                 row = [
                     time,
                     *position,
                     node_id,
                     parent_id,
                     track_id,
+                    lineage_id,
+                    color,
                 ]
                 f.write("\n")
                 f.write(",".join(map(str, row)))

src/motile_tracker/data_views/views/layers/track_graph.py

@@ -89,11 +89,14 @@ def __init__(
             name=name,
             tail_length=3,
             color_by="track_id",
+            blending="translucent_no_depth",
         )
 
         self.colormaps_dict["track_id"] = self.tracks_viewer.colormap
         self.tracks_layer_graph = copy.deepcopy(self.graph)  # for restoring graph later
         self.colormap = "turbo"  # just to 'refresh' the track_id colormap, we do not actually use turbo
+        self.visible_plane_tracks = "all"
+        self.visible_tracks = "all"
 
     def _refresh(self):
         """Refreshes the displayed tracks based on the graph in the current tracks_viewer.tracks"""
@@ -108,13 +111,50 @@ def _refresh(self):
         self.colormaps_dict["track_id"] = self.tracks_viewer.colormap
         self.colormap = "turbo"  # just to 'refresh' the track_id colormap, we do not actually use turbo
 
-    def update_track_visibility(self, visible: list[int] | str) -> None:
+    def update_track_visibility(
+        self,
+        visible: list[int] | str | None = None,
+        plane_nodes: list[int] | str | None = None,
+    ) -> None:
         """Optionally show only the tracks of a current lineage"""
 
-        if visible == "all":
+        if (
+            self.tracks_viewer.viewer.dims.ndisplay == 2
+            and self.tracks_viewer.viewer.dims.ndim == 4
+        ):
+            # everything should be invisible for 3D data viewed in 2D, because tracks are not shown at the correct slices.
+            visible = []
+
+        else:
+            if visible is not None:
+                self.visible_tracks = visible
+            if plane_nodes is not None:
+                self.visible_plane_tracks = plane_nodes
+
+            if isinstance(self.visible_tracks, str) and isinstance(
+                self.visible_plane_tracks, str
+            ):
+                visible = "all"
+            elif not isinstance(self.visible_tracks, str) and isinstance(
+                self.visible_plane_tracks, str
+            ):
+                visible = self.visible_tracks
+            elif isinstance(self.visible_tracks, str) and not isinstance(
+                self.visible_plane_tracks, str
+            ):
+                visible = self.visible_plane_tracks
+            else:
+                visible = list(
+                    set(self.visible_tracks).intersection(
+                        set(self.visible_plane_tracks)
+                    )
+                )
+
+        if isinstance(visible, str):
             self.track_colors[:, 3] = 1
             self.graph = self.tracks_layer_graph
         else:
+            visible = [self.tracks_viewer.tracks.get_track_id(node) for node in visible]
             track_id_mask = np.isin(
                 self.properties["track_id"],
                 visible,
@@ -131,3 +171,5 @@ def update_track_visibility(self, visible: list[int] | str) -> None:
                 self.display_graph = False  # empty dicts to not trigger update (bug?) so disable the graph entirely as a workaround
             else:
                 self.display_graph = True
+
+        self.events.rebuild_tracks()  # fire the event to update the colors

src/motile_tracker/data_views/views/layers/track_labels.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import random
+import time
 from typing import TYPE_CHECKING
 
 import napari
@@ -110,33 +111,8 @@ def __init__(
         self.bind_key("z")(self.tracks_viewer.undo)
         self.bind_key("r")(self.tracks_viewer.redo)
 
-        # Connect click events to node selection
-        @self.mouse_drag_callbacks.append
-        def click(_, event):
-            if (
-                event.type == "mouse_press"
-                and self.mode == "pan_zoom"
-                and not (
-                    self.tracks_viewer.mode == "lineage"
-                    and self.viewer.dims.ndisplay == 3
-                )
-            ):  # disable selecting in lineage mode in 3D
-                label = self.get_value(
-                    event.position,
-                    view_direction=event.view_direction,
-                    dims_displayed=event.dims_displayed,
-                    world=True,
-                )
-
-                if (
-                    label is not None
-                    and label != 0
-                    and self.colormap.map(label)[-1] != 0
-                ):  # check opacity (=visibility) in the colormap
-                    append = "Shift" in event.modifiers
-                    self.tracks_viewer.selected_nodes.add(label, append)
-
         # Listen to paint events and changing the selected label
+        self.mouse_drag_callbacks.append(self.click)
         self.events.paint.connect(self._on_paint)
         self.tracks_viewer.selected_nodes.list_updated.connect(
             self.update_selected_label
@@ -145,6 +121,44 @@ def click(_, event):
         self.events.mode.connect(self._check_mode)
         self.viewer.dims.events.current_step.connect(self._ensure_valid_label)
 
+    # Connect click events to node selection
+    def click(self, _, event):
+        if (
+            event.type == "mouse_press"
+            and self.mode == "pan_zoom"
+            and not (
+                self.tracks_viewer.mode == "lineage" and self.viewer.dims.ndisplay == 3
+            )
+        ):  # disable selecting in lineage mode in 3D
+            # differentiate between click and drag
+            mouse_press_time = time.time()
+            dragged = False
+            yield
+            # on move
+            while event.type == "mouse_move":
+                dragged = True
+                yield
+            if dragged and time.time() - mouse_press_time < 0.5:
+                dragged = False  # suppress micro drag events and treat them as click
+            # on release
+            if not dragged:
+                label = self.get_value(
+                    event.position,
+                    view_direction=event.view_direction,
+                    dims_displayed=event.dims_displayed,
+                    world=True,
+                )
+                self.process_click(event, label)
+
+    def process_click(self, event: Event, label: int):
+        if (
+            label is not None and label != 0 and self.colormap.map(label)[-1] != 0
+        ):  # check opacity (=visibility) in the colormap
+            append = "Shift" in event.modifiers
+            self.tracks_viewer.selected_nodes.add(label, append)
+        else:
+            self.tracks_viewer.selected_nodes.reset()
+
     def _get_colormap(self) -> DirectLabelColormap:
         """Get a DirectLabelColormap that maps node ids to their track ids, and then
         uses the tracks_viewer.colormap to map from track_id to color.
@@ -225,8 +239,8 @@ def _parse_paint_event(self, event_val):
             mask = concatenated_values == old_value
             indices = tuple(concatenated_indices[dim][mask] for dim in range(ndim))
             time_points = np.unique(indices[0])
-            for time in time_points:
-                time_mask = indices[0] == time
+            for t in time_points:
+                time_mask = indices[0] == t
                 actions.append(
                     (tuple(indices[dim][time_mask] for dim in range(ndim)), old_value)
                 )