Merge branch 'main' of https://github.com/Hou-Lab-CSHL/cheese3d

Author: lzhan132

README.md

@@ -4,7 +4,7 @@
 [![Data](https://img.shields.io/badge/download-demodata-blue)](https://labshare.cshl.edu/shares/houlab/www-data/cheese3d_paper_data/cheese3d_demo.tar.gz)
 <!--[Download demo data](https://labshare.cshl.edu/shares/houlab/www-data/cheese3d_paper_data/cheese3d_demo.tar.gz)-->
 
-Cheese3D is a pipeline for tracking mouse facial movements built on top of existing tools like [DeepLabCut](https://github.com/DeepLabCut/DeepLabCut) and [Anipose](https://github.com/lambdaloop/anipose). By tracking anatomically-informed keypoints using multiple cameras registered in 3D, our pipeline produces sensitive, high-precision facial movement data that can be related internal state (e.g., electrophysiology).
+Cheese3D is a pipeline for tracking mouse facial movements built on top of existing tools ([DeepLabCut](https://github.com/DeepLabCut/DeepLabCut) and [Anipose](https://github.com/lambdaloop/anipose)). By tracking anatomically-informed keypoints using multiple cameras registered in 3D, our pipeline produces sensitive, high-precision facial movement data that can be related internal state (e.g., electrophysiology).
 
 <p align="center">
    <img src="docs/source/_static/Cheese3D.gif" alt="Animation of Cheese3D pipeline", width=60%>

docs/source/index.rst

@@ -3,7 +3,7 @@
 Cheese3D
 ========
 
-Cheese3D is a pipeline for tracking mouse facial movements built on top of existing tools like https://github.com/DeepLabCut/DeepLabCut and https://github.com/lambdaloop/anipose. By tracking anatomically-informed keypoints using multiple cameras registered in 3D, our pipeline produces sensitive, high-precision facial movement data that can be related internal state (e.g., electrophysiology).
+Cheese3D is a pipeline for tracking mouse facial movements built on top of existing tools (https://github.com/DeepLabCut/DeepLabCut and https://github.com/lambdaloop/anipose). By tracking anatomically-informed keypoints using multiple cameras registered in 3D, our pipeline produces sensitive, high-precision facial movement data that can be related internal state (e.g., electrophysiology).
 
 .. image:: /_static/Cheese3D.gif
     :width: 59%

packages/cheese3d-annotator/README.md

@@ -111,3 +111,55 @@ def ensure_images_in_yaml(image_files: List[str],
     with open(yaml_path, "w") as f:
         yaml.safe_dump(annotations, f)
     print(f"▶︎ Added {n_files_added} new image(s) to {os.path.basename(yaml_path)}")
+
+def save_shapes_yaml(layers_by_part, filenames, path, labeler=""):
+    """
+    Save shape annotations using image filenames as keys.
+
+    Parameters
+    ----------
+    layers_by_part : dict
+        Dict of {part: napari shape layer}
+    filenames : list of str
+        List of frame image paths
+    path : str
+        Path to save annotation.yaml
+    labeler : str
+        Labeler name to include
+    """
+    out = {}
+    for part, layer in layers_by_part.items():
+        part_data = {}
+        for shape in layer.data:
+            shape = np.asarray(shape)
+            if shape.ndim != 2 or shape.shape[1] != 3 or np.isnan(shape).any():
+                continue
+            z = int(round(shape[0, 0]))
+            if z < 0 or z >= len(filenames):
+                continue
+            fname = os.path.basename(filenames[z])
+            part_data.setdefault(fname, []).append(shape[:, 1:].tolist())  # strip z-axis
+        out[part] = part_data
+
+    with open(path, "w") as f:
+        yaml.safe_dump(out, f)
+
+
+def load_shapes_yaml(path, filenames):
+    with open(path, "r") as f:
+        raw = yaml.safe_load(f) or {}
+
+    file_to_index = {os.path.basename(f): i for i, f in enumerate(filenames)}
+
+    out = {}
+    for part, part_data in raw.items():
+        data = []
+        for fname, shapes in part_data.items():
+            z = file_to_index.get(fname, None)
+            if z is None:
+                continue
+            for s in shapes:
+                shape = np.array([[z, y, x] for y, x in s], dtype=np.float32)
+                data.append(shape)
+        out[part] = data
+    return out

packages/cheese3d-annotator/cheese3d_annotator/data.py

@@ -3,13 +3,18 @@ display_name: Cheese3D Annotator
 contributions:
   commands:
     - id: cheese3d-annotator.frame_picker
-      python_name: cheese3d-annotator.widget:FramePickerWidget
+      python_name: cheese3d_annotator.widget:FramePickerWidget
       title: Cheese3D Frame Picker
     - id: cheese3d-annotator.frame_annotator
-      python_name: cheese3d-annotator.widget:FrameAnnotatorWidget
+      python_name: cheese3d_annotator.widget:FrameAnnotatorWidget
       title: Cheese3D Frame Annotator
+    - id: cheese3d-annotator.curve_annotator
+      python_name: cheese3d_annotator.widget:CurveAnnotatorWidget
+      title: Cheese3D Curve Annotator (experimental)
   widgets:
     - command: cheese3d-annotator.frame_annotator
       display_name: Cheese3D Frame Annotator
     - command: cheese3d-annotator.frame_picker
       display_name: Cheese3D Frame Picker
+    - command: cheese3d-annotator.curve_annotator
+      display_name: Cheese3D Curve Annotator (experimental)

packages/cheese3d-annotator/cheese3d_annotator/napari.yaml

@@ -4,7 +4,7 @@
 import numpy as np
 import pandas as pd
 from typing import List
-from magicgui.widgets import Container, FileEdit, ComboBox, Label, CheckBox, PushButton
+from magicgui.widgets import Container, FileEdit, LineEdit, ComboBox, Label, CheckBox, PushButton
 from qtpy.QtWidgets import QListWidget, QListWidgetItem, QMessageBox, QSizePolicy
 from qtpy.QtGui import QFont, QImage, QPixmap, QIcon
 from qtpy.QtCore import QSize
@@ -20,7 +20,9 @@
                                      write_annotations,
                                      create_empty_annotations,
                                      find_keypoint_conflicts,
-                                     ensure_images_in_yaml)
+                                     ensure_images_in_yaml,
+                                     save_shapes_yaml,
+                                     load_shapes_yaml)
 
 class FrameAnnotatorWidget(Container):
     def __init__(self, viewer: Viewer):
@@ -565,3 +567,117 @@ def jump_to_time(self, item):
             self.viewer.dims.current_step = (frame_index,) + self.viewer.dims.current_step[1:]
         except Exception as e:
             QMessageBox.warning(None, "Jump Failed", str(e))
+
+class CurveAnnotatorWidget(Container):
+    def __init__(self, viewer: Viewer):
+        super().__init__()
+        self.viewer = viewer
+        self.viewer.layers.clear()
+        self.viewer.window.remove_dock_widget('all')
+        self.viewer.grid.enabled = False
+
+        self.root_folder = FileEdit(label="Root Folder", mode="d")
+        self.labeler_name = LineEdit(label="Labeler", value="houlab")
+
+        self.folder_list = QListWidget()
+        self.folder_list.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Minimum)
+        self.folder_list.setMaximumHeight(500)
+        self.folder_list.setFont(QFont("", 12))
+        self.folder_list.setStyleSheet("QListWidget::item { height: 32px; }")
+        self.folder_list.itemClicked.connect(self.load_subfolder)
+
+        self.help_label = Label(value="""
+        <b>Instructions:</b><br>
+        - Select Root Folder.<br>
+        - Pick a frame folder to annotate.<br>
+        - Use shape tool to draw curves on each frame.<br>
+        - Use slider to move between frames.<br>
+        - Press <b>S</b> to save.
+        """)
+
+        self.extend([
+            self.labeler_name,
+            self.root_folder,
+        ])
+        self.native.layout().addWidget(self.folder_list)
+        self.native.layout().addWidget(self.help_label.native)
+
+        self.root_folder.changed.connect(self.refresh_folders)
+
+        self.parts = ["right_ear", "left_ear"]
+        self.shape_layers = {}  # {part: shape_layer}
+        self.current_folder = None
+        self.last_frame = None
+
+        self._bind_keys()
+        # TODO: add on change callback for parts_dropdown
+        self.viewer.dims.events.current_step.connect(self._on_frame_change)
+
+    def refresh_folders(self):
+        self.folder_list.clear()
+        base = self.root_folder.value
+        if not base or not os.path.isdir(base):
+            return
+        for f in sorted(os.listdir(base)):
+            if os.path.isdir(os.path.join(base, f)):
+                self.folder_list.addItem(f)
+
+    def _on_frame_change(self, event):
+        frame = self.viewer.dims.current_step[0]
+        if self.last_frame is not None and frame != self.last_frame:
+            self._save(self.viewer)  # auto-save when frame changes
+        self.last_frame = frame
+
+    def _bind_keys(self):
+        @self.viewer.bind_key("y", overwrite=True)
+        def _save(viewer):
+            self._save(viewer)
+
+    def load_subfolder(self, item: QListWidgetItem):
+        if self.current_folder and self.shape_layers:
+            self._save(self.viewer)
+
+        folder_name = item.text()
+        folder_path = os.path.join(self.root_folder.value, folder_name)
+        self.viewer.layers.clear()
+        self.shape_layers.clear()
+
+        # Load image stack
+        self.filenames = sorted(glob(os.path.join(folder_path, "*.png")))
+        if not self.filenames:
+            QMessageBox.warning(None, "No images", f"No PNGs found in {folder_name}")
+            return
+
+        stack = np.stack([
+            rgb2gray(imread(p)) if imread(p).ndim == 3 else imread(p)
+            for p in self.filenames
+        ])
+        self.viewer.add_image(stack, name=folder_name)
+        self.viewer.dims.ndisplay = 2
+        self.viewer.dims.axis_labels = ("frame", "y", "x")
+        self.current_folder = folder_path
+        self.last_frame = self.viewer.dims.current_step[0]
+
+        # Load shapes
+        annotation_path = os.path.join(folder_path, "annotation.yaml")
+        all_data = load_shapes_yaml(annotation_path, self.filenames) if os.path.exists(annotation_path) else {}
+
+        for part in self.parts:
+            layer = self.viewer.add_shapes(
+                name=f"{part}_shapes", shape_type="path", ndim=3, edge_width=4
+            )
+            if part in all_data and all_data[part]:
+                layer.data = all_data[part]
+                layer.shape_type = "path"
+            self.shape_layers[part] = layer
+
+    def _save(self, viewer):
+        if self.current_folder and self.shape_layers:
+            labeler = self.labeler_name.value
+            save_shapes_yaml(
+                self.shape_layers,
+                self.filenames,
+                os.path.join(self.current_folder, "annotation.yaml"),
+                labeler=labeler
+            )
+            print(f"✅ Saved to {self.current_folder}/annotation.yaml")

packages/cheese3d-annotator/cheese3d_annotator/widget.py

@@ -7,8 +7,8 @@ authors = [
     { name = "Kyle Daruwalla", email = "daruwal@cshl.edu"},
     { name = "Rubin Zhao", email = "rzhao@cshl.edu"}
 ]
-readme = "README.md"
-requires-python = "==3.10.*"
+readme = "../../README.md"
+requires-python = ">=3.10"
 dependencies = [
     "tables>=3.7.0",
     "pandas>=2.1.4",

packages/cheese3d-annotator/pyproject.toml

@@ -26,6 +26,6 @@ def extract_frames(self, videos: Optional[List[Path]] = None):
         """Extract frames from videos."""
         raise NotImplementedError("This method should be implemented by subclasses.")
 
-    def train(self, gpu):
+    def train(self, gpu, iterate_dataset: bool = True):
         """Train the model using GPU ID `gpu`."""
         raise NotImplementedError("This method should be implemented by subclasses.")