Merge pull request #95 from funkelab/91-implement-import-dialog-for-geff-graph

91 implement import dialog for geff graph

Author: cmalinmayor

finn/track_data_views/views_coordinator/tracks_list.py

@@ -28,6 +28,9 @@
 from finn.track_import_export.menus.import_external_tracks_dialog import (
     ImportTracksDialog,
 )
+from finn.track_import_export.menus.import_from_geff.geff_import_dialog import (
+    ImportGeffDialog,
+)
 
 
 class TrackListWidget(QWidget):
@@ -110,7 +113,9 @@ def __init__(self):
 
         load_menu = QHBoxLayout()
         self.dropdown_menu = QComboBox()
-        self.dropdown_menu.addItems(["Motile Run", "External tracks from CSV"])
+        self.dropdown_menu.addItems(
+            ["Motile Run", "External tracks from CSV", "External tracks from geff"]
+        )
 
         load_button = QPushButton("Load")
         load_button.clicked.connect(self.load_tracks)
@@ -123,6 +128,14 @@ def __init__(self):
         layout.addLayout(load_menu)
         self.setLayout(layout)
 
+    def _load_tracks_from_geff(self):
+        dialog = ImportGeffDialog()
+        if dialog.exec_() == QDialog.Accepted:
+            tracks = dialog.tracks
+            name = dialog.name
+            if tracks is not None:
+                self.add_tracks(tracks, name, select=True)
+
     def _load_external_tracks(self):
         dialog = ImportTracksDialog()
         if dialog.exec_() == QDialog.Accepted:
@@ -251,6 +264,8 @@ def load_tracks(self):
             self.load_motile_run()
         elif self.dropdown_menu.currentText() == "External tracks from CSV":
             self._load_external_tracks()
+        elif self.dropdown_menu.currentText() == "External tracks from geff":
+            self._load_tracks_from_geff()
 
     def load_motile_run(self):
         """Load a set of tracks from disk. The user selects the directory created

finn/track_import_export/menus/import_from_geff/__init__.py

@@ -0,0 +1,206 @@
+from pathlib import Path
+
+from funtracks.import_export.import_from_geff import import_from_geff
+from qtpy.QtCore import Qt
+from qtpy.QtWidgets import (
+    QApplication,
+    QDialog,
+    QHBoxLayout,
+    QMessageBox,
+    QPushButton,
+    QScrollArea,
+    QSizePolicy,
+    QVBoxLayout,
+    QWidget,
+)
+
+from finn.track_import_export.menus.import_from_geff.geff_import_widget import (
+    ImportGeffWidget,
+)
+from finn.track_import_export.menus.import_from_geff.geff_prop_map_widget import (
+    StandardFieldMapWidget,
+)
+from finn.track_import_export.menus.import_from_geff.geff_scale_widget import ScaleWidget
+from finn.track_import_export.menus.import_from_geff.geff_segmentation_widgets import (
+    SegmentationWidget,
+)
+
+
+class ImportGeffDialog(QDialog):
+    """Dialgo for importing external tracks from a geff file"""
+
+    def __init__(self):
+        super().__init__()
+        self.setWindowTitle("Import external tracks from geff")
+        self.name = "Tracks from Geff"
+
+        # cancel and finish buttons
+        self.button_layout = QHBoxLayout()
+        self.cancel_button = QPushButton("Cancel")
+        self.finish_button = QPushButton("Finish")
+        self.finish_button.setEnabled(False)
+        self.button_layout.addWidget(self.cancel_button)
+        self.button_layout.addWidget(self.finish_button)
+
+        # Connect button signals
+        self.cancel_button.clicked.connect(self._cancel)
+        self.finish_button.clicked.connect(self._finish)
+        self.cancel_button.setDefault(False)
+        self.cancel_button.setAutoDefault(False)
+        self.finish_button.setDefault(False)
+        self.finish_button.setAutoDefault(False)
+
+        # Initialize widgets and connect to update signals
+        self.geff_widget = ImportGeffWidget()
+        self.geff_widget.update_buttons.connect(self._update_segmentation_widget)
+        self.segmentation_widget = SegmentationWidget(root=self.geff_widget.root)
+        self.segmentation_widget.none_radio.toggled.connect(
+            self._toggle_scale_widget_and_seg_id
+        )
+        self.segmentation_widget.segmentation_widget.seg_path_updated.connect(
+            self._update_finish_button
+        )
+        self.prop_map_widget = StandardFieldMapWidget()
+        self.geff_widget.update_buttons.connect(self._update_field_map_widget)
+        self.scale_widget = ScaleWidget()
+
+        self.content_widget = QWidget()
+        main_layout = QVBoxLayout(self.content_widget)
+        main_layout.addWidget(self.geff_widget)
+        main_layout.addWidget(self.segmentation_widget)
+        main_layout.addWidget(self.prop_map_widget)
+        main_layout.addWidget(self.scale_widget)
+        main_layout.addLayout(self.button_layout)
+        self.content_widget.setLayout(main_layout)
+        self.content_widget.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Maximum)
+
+        # Create scroll area
+        self.scroll_area = QScrollArea()
+        self.scroll_area.setWidget(self.content_widget)
+        self.scroll_area.setWidgetResizable(True)
+        self.scroll_area.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
+        self.scroll_area.setMinimumWidth(500)
+        self.scroll_area.setSizePolicy(
+            QSizePolicy.Preferred, QSizePolicy.MinimumExpanding
+        )
+        self.content_widget.setSizePolicy(QSizePolicy.Preferred, QSizePolicy.Maximum)
+
+        dialog_layout = QVBoxLayout()
+        dialog_layout.addWidget(self.scroll_area)
+        self.setLayout(dialog_layout)
+        self.setSizePolicy(self.sizePolicy().horizontalPolicy(), QSizePolicy.Minimum)
+
+    def _resize_dialog(self):
+        """Dynamic widget resizing depending on the visible contents"""
+
+        self.content_widget.adjustSize()
+        self.content_widget.updateGeometry()
+
+        content_hint = self.content_widget.sizeHint()
+        screen_geometry = QApplication.primaryScreen().availableGeometry()
+
+        max_height = int(screen_geometry.height() * 0.85)
+        new_height = min(content_hint.height(), max_height)
+        new_width = max(content_hint.width(), 500)
+
+        self.resize(new_width, new_height)
+
+        # Center horizontally, but upwards if too tall
+        screen_center = screen_geometry.center()
+        x = screen_center.x() - self.width() // 2
+
+        if new_height < screen_geometry.height():
+            y = screen_center.y() - new_height // 2
+        else:
+            y = screen_geometry.top() + 50
+
+        self.move(x, y)
+
+    def _update_segmentation_widget(self) -> None:
+        """Refresh the segmentation widget based on the geff root group."""
+
+        if self.geff_widget.root is not None:
+            self.segmentation_widget.update_root(self.geff_widget.root)
+        else:
+            self.segmentation_widget.setVisible(False)
+        self._update_finish_button()
+        self._resize_dialog()
+
+    def _update_field_map_widget(self) -> None:
+        """Prefill the field map widget with the geff metadata and graph attributes."""
+
+        if self.geff_widget.root is not None:
+            self.prop_map_widget.update_mapping(
+                self.geff_widget.root, self.segmentation_widget.include_seg()
+            )
+
+            self.scale_widget._prefill_from_metadata(
+                dict(self.geff_widget.root.attrs.get("geff", {}))
+            )
+            self.scale_widget.setVisible(self.segmentation_widget.include_seg())
+        else:
+            self.prop_map_widget.setVisible(False)
+            self.scale_widget.setVisible(False)
+
+        self._update_finish_button()
+        self._resize_dialog()
+
+    def _update_finish_button(self):
+        """Update the finish button status depending on whether a segmentation is required
+        and whether a valid geff root is present."""
+
+        include_seg = self.segmentation_widget.include_seg()
+        has_seg = self.segmentation_widget.get_segmentation() is not None
+        valid_seg = not (include_seg and not has_seg)
+        self.finish_button.setEnabled(self.geff_widget.root is not None and valid_seg)
+
+    def _toggle_scale_widget_and_seg_id(self, checked: bool) -> None:
+        """Toggle visibility of the scale widget based on the 'None' radio button state,
+        and update the visibility of the 'seg_id' combobox in the prop map widget."""
+
+        self.scale_widget.setVisible(not checked)
+
+        # Also remove the seg_id from the fields widget
+        if len(self.prop_map_widget.mapping_widgets) > 0:
+            self.prop_map_widget.mapping_widgets["seg_id"].setVisible(not checked)
+            self.prop_map_widget.mapping_labels["seg_id"].setVisible(not checked)
+            self.prop_map_widget.optional_features["area"]["recompute"].setEnabled(
+                not checked
+            )
+
+        self._update_finish_button()
+        self._resize_dialog()
+
+    def _cancel(self) -> None:
+        """Close the dialog without loading tracks."""
+        self.reject()
+
+    def _finish(self) -> None:
+        """Tries to read the geff file and optional segmentation image and apply the
+        attribute to column mapping to construct a Tracks object"""
+
+        if self.geff_widget.root is not None:
+            store_path = Path(self.geff_widget.root.store.path)  # e.g. /.../my_store.zarr
+            group_path = Path(self.geff_widget.root.path)  # e.g. 'tracks'
+            geff_dir = store_path / group_path
+
+            self.name = self.geff_widget.dir_name
+            scale = self.scale_widget.get_scale()
+
+            segmentation = self.segmentation_widget.get_segmentation()
+            name_map = self.prop_map_widget.get_name_map()
+            name_map = {k: (None if v == "None" else v) for k, v in name_map.items()}
+            extra_features = self.prop_map_widget.get_optional_props()
+
+            try:
+                self.tracks = import_from_geff(
+                    geff_dir,
+                    name_map,
+                    segmentation_path=segmentation,
+                    extra_features=extra_features,
+                    scale=scale,
+                )
+            except (ValueError, OSError, FileNotFoundError, AssertionError) as e:
+                QMessageBox.critical(self, "Error", f"Failed to load tracks: {e}")
+                return
+            self.accept()

finn/track_import_export/menus/import_from_geff/geff_import_dialog.py

@@ -0,0 +1,40 @@
+import zarr
+from qtpy.QtWidgets import (
+    QLayout,
+)
+
+
+def find_geff_group(group: zarr.Group) -> zarr.Group | None:
+    """Recursively search for a Zarr group with 'geff' in its .zattrs.
+    Args:
+        group (zarr.Group): The Zarr group to search within.
+    Returns:
+        zarr.Group | None: The first group found with 'geff' in its .zattrs, or None if
+        not found.
+    """
+
+    if "geff" in group.attrs:
+        return group
+
+    for key in group.group_keys():
+        subgroup = group[key]
+        if isinstance(subgroup, zarr.Group):
+            found = find_geff_group(subgroup)
+            if found:
+                return found
+    return None
+
+
+def clear_layout(layout: QLayout) -> None:
+    """Recursively clear all widgets and layouts in a QLayout.
+    Args:
+        layout (QLayout): The layout to clear.
+    """
+    while layout.count():
+        item = layout.takeAt(0)
+        widget = item.widget()
+        if widget is not None:
+            widget.setParent(None)
+        # If the item is a layout itself, clear it recursively
+        elif item.layout() is not None:
+            clear_layout(item.layout())

finn/track_import_export/menus/import_from_geff/geff_import_utils.py

@@ -0,0 +1,102 @@
+import os
+from pathlib import Path
+
+import zarr
+from psygnal import Signal
+from qtpy.QtCore import Qt
+from qtpy.QtWidgets import (
+    QFileDialog,
+    QGroupBox,
+    QHBoxLayout,
+    QLineEdit,
+    QMessageBox,
+    QPushButton,
+    QVBoxLayout,
+    QWidget,
+)
+from zarr.storage import FSStore
+
+from finn.track_import_export.menus.import_from_geff.geff_import_utils import (
+    find_geff_group,
+)
+
+
+class ImportGeffWidget(QWidget):
+    """QWidget for selecting geff directory"""
+
+    update_buttons = Signal()
+
+    def __init__(self):
+        super().__init__()
+
+        self.root = None
+        self.dir_name = None
+
+        # QlineEdit for geff file path and browse button
+        self.geff_path_line = QLineEdit(self)
+        self.geff_path_line.setFocus()
+        self.geff_path_line.setFocusPolicy(Qt.StrongFocus)
+        self.geff_path_line.returnPressed.connect(self._on_line_editing_finished)
+        self.geff_browse_button = QPushButton("Browse", self)
+        self.geff_browse_button.setAutoDefault(0)
+        self.geff_browse_button.clicked.connect(self._browse_geff)
+
+        browse_layout = QHBoxLayout()
+        browse_layout.addWidget(self.geff_path_line)
+        browse_layout.addWidget(self.geff_browse_button)
+
+        box = QGroupBox("Path to geff zarr directory")
+        box_layout = QVBoxLayout()
+        box_layout.addLayout(browse_layout)
+        box.setLayout(box_layout)
+        main_layout = QVBoxLayout()
+        main_layout.addWidget(box)
+        self.setLayout(main_layout)
+
+    def _on_line_editing_finished(self) -> None:
+        """Load the geff group when the user presses Enter in the path line"""
+
+        folder_path = self.geff_path_line.text().strip()
+        if not folder_path:
+            self.root = None
+            self.update_buttons.emit()  # to remove any widgets and disable finish button
+        else:
+            # try to load, will raise an error if no zarr group can be found
+            self._load_geff(Path(folder_path))
+
+    def _browse_geff(self) -> None:
+        """Open File dialog to select geff folder"""
+
+        folder = QFileDialog.getExistingDirectory(self, "Select Geff Zarr directory")
+        if folder:
+            folder_path = Path(folder)
+            self.geff_path_line.setText(str(folder_path))
+            self._load_geff(folder_path)
+
+    def _load_geff(self, folder_path: Path) -> None:
+        """Load the graph and display the geffFieldMapWidget"""
+
+        self.root = None
+        if not os.path.exists(folder_path):
+            QMessageBox.critical(self, "Error", f"Path does not exist: {folder_path}")
+            self.update_buttons.emit()
+            return
+        try:
+            store = FSStore(folder_path)
+            root = zarr.group(store=store)
+        except (KeyError, ValueError, OSError) as e:
+            QMessageBox.critical(self, "Error", f"Could not open zarr store: {e}")
+            self.update_buttons.emit()
+            return
+
+        geff_group = find_geff_group(root)
+        if geff_group is None:
+            QMessageBox.critical(
+                self, "Error", "No geff group found in the selected directory."
+            )
+            self.update_buttons.emit()
+            return
+
+        self.root = geff_group
+        self.dir_name = os.path.basename(folder_path)
+        self.update_buttons.emit()