feat: tracking of sarcomere vectors by computation of optical flow, position advection, and vector matching

Author: danihae

sarcasm/structure.py

@@ -40,7 +40,7 @@
     kymograph,
     detection,
     loi_detection,
-    # sarcomere_tracking
+    sarcomere_tracking,
 )
 
 class Structure(SarcAsM):
@@ -1615,188 +1615,188 @@ def delete_lois(self):
             except Exception as e:
                 logger.debug(f"Error deleting LOI directory: {e}. Continuing anyway.")
 
-    # def track_sarcomere_vectors(
-    #     self,
-    #     frames: Union[str, int, List[int], np.ndarray] = 'all',
-    #     max_distance: float = 0.5,
-    #     memory: int = 5,
-    #     min_track_length: int = 10,
-    #     flow_blur_sigma: Optional[float] = None,
-    #     use_hungarian: bool = True,
-    #     length_weight: float = 2.0,
-    #     return_flow: bool = False,
-    #     progress_notifier: ProgressNotifier = ProgressNotifier.progress_notifier_tqdm()
-    # ) -> Optional[List[np.ndarray]]:
-    #     """
-    #     Track sarcomere vectors across frames using optical flow.
-
-    #     This method uses dense optical flow computed on Z-band and M-band masks
-    #     to predict vector motion between frames, then matches predicted positions
-    #     to detected vectors using greedy exclusive assignment.
-
-    #     Parameters
-    #     ----------
-    #     frames : {'all', int, list, np.ndarray}, optional
-    #         Frames for tracking ('all' for all frames, int for a single frame,
-    #         list or ndarray for selected frames). Defaults to 'all'.
-    #     max_distance : float, optional
-    #         Maximum matching distance in µm. Vectors predicted to be further
-    #         than this from any detection are considered lost. Default 0.5 µm.
-    #     memory : int, optional
-    #         Number of frames to keep tracks "sleeping" before termination.
-    #         Sleeping tracks can be re-linked if a matching vector reappears.
-    #         Default 5 frames.
-    #     min_track_length : int, optional
-    #         Minimum track length to keep. Shorter tracks are filtered out
-    #         and their vectors marked as untracked (track_id = -1). Default 10.
-    #     flow_blur_sigma : float, optional
-    #         Gaussian blur sigma (in pixels) to apply to masks before optical
-    #         flow computation. Can help with noisy masks. Default None (no blur).
-    #     use_hungarian : bool, optional
-    #         If True, use Hungarian algorithm for optimal matching. This gives
-    #         better results when vectors are densely packed. Default True.
-    #     length_weight : float, optional
-    #         Weight for sarcomere length difference in matching cost when using
-    #         Hungarian algorithm. Higher values favor matching vectors with
-    #         similar sarcomere lengths. Default 2.0.
-    #     return_flow : bool, optional
-    #         If True, also return the computed optical flow fields.
-    #         Default False.
-    #     progress_notifier : ProgressNotifier, optional
-    #         Progress notifier for tracking progress.
-
-    #     Returns
-    #     -------
-    #     flow_fields : list of np.ndarray, optional
-    #         Only returned if return_flow=True. List of flow fields,
-    #         each of shape (H, W, 2).
-
-    #     Notes
-    #     -----
-    #     Results are stored in `self.data` with the following keys:
-
-    #     - ``'track_ids'``: list of (N_t,) arrays with track ID for each vector
-    #       per frame. Track ID of -1 indicates untracked vectors.
-    #     - ``'n_tracks'``: total number of unique trajectories.
-    #     - ``'track_lengths'``: (n_tracks,) array with number of frames each
-    #       track spans.
-    #     - ``'params.track_sarcomere_vectors.*'``: tracking parameters used.
-
-    #     Requires :meth:`analyze_sarcomere_vectors` to be run first.
-
-    #     Examples
-    #     --------
-    #     >>> structure.analyze_sarcomere_vectors(frames='all')
-    #     >>> structure.track_sarcomere_vectors(max_distance=0.5, memory=5)
-    #     >>> track_ids = structure.data['track_ids']
-    #     >>> # track_ids[frame][vector_idx] gives the track ID for that vector
-    #     """
-    #     # Validate prerequisites
-    #     if 'pos_vectors_px' not in self.data:
-    #         raise ValueError(
-    #             "Sarcomere vectors not yet analyzed. Run analyze_sarcomere_vectors first."
-    #         )
-
-    #     if not os.path.exists(self.file_zbands) or not os.path.exists(self.file_mbands):
-    #         raise FileNotFoundError(
-    #             "Z-band or M-band mask not found. Please run detect_sarcomeres first."
-    #         )
-
-    #     # Determine frames to track
-    #     analyzed_frames = self.data.get('params.analyze_sarcomere_vectors.frames', [])
-    #     if isinstance(frames, str) and frames == 'all':
-    #         list_frames = list(range(self.metadata.n_stack))
-    #     elif np.issubdtype(type(frames), np.integer):
-    #         list_frames = [frames]
-    #     elif isinstance(frames, (list, np.ndarray)):
-    #         list_frames = [int(f) for f in frames]
-    #     else:
-    #         raise ValueError('frames argument not valid')
-
-    #     # Validate frames have been analyzed
-    #     if not set(list_frames).issubset(set(analyzed_frames)):
-    #         missing = set(list_frames) - set(analyzed_frames)
-    #         raise ValueError(
-    #             f"Frames {missing} have not been analyzed. "
-    #             f"Run analyze_sarcomere_vectors first for these frames."
-    #         )
-
-    #     logger.info(f"Starting sarcomere vector tracking for {len(list_frames)} frames...")
-
-    #     # Load masks
-    #     zbands_stack = tifffile.imread(self.file_zbands)
-    #     mbands_stack = tifffile.imread(self.file_mbands)
-
-    #     # Ensure 3D
-    #     if zbands_stack.ndim == 2:
-    #         zbands_stack = np.expand_dims(zbands_stack, axis=0)
-    #     if mbands_stack.ndim == 2:
-    #         mbands_stack = np.expand_dims(mbands_stack, axis=0)
-
-    #     # Select frames
-    #     zbands_stack = zbands_stack[list_frames]
-    #     mbands_stack = mbands_stack[list_frames]
-
-    #     # Prepare vector data for selected frames
-    #     positions_per_frame = [self.data['pos_vectors_px'][f] for f in list_frames]
-    #     lengths_per_frame = [self.data['sarcomere_length_vectors'][f] for f in list_frames]
-    #     orientations_per_frame = [self.data['sarcomere_orientation_vectors'][f] for f in list_frames]
-
-    #     # Run tracking
-    #     result = sarcomere_tracking.track_sarcomere_vectors(
-    #         zbands_stack=zbands_stack,
-    #         mbands_stack=mbands_stack,
-    #         positions_per_frame=positions_per_frame,
-    #         sarcomere_lengths_per_frame=lengths_per_frame,
-    #         orientations_per_frame=orientations_per_frame,
-    #         pixelsize=self.metadata.pixelsize,
-    #         max_distance=max_distance,
-    #         memory=memory,
-    #         min_track_length=min_track_length,
-    #         flow_blur_sigma=flow_blur_sigma,
-    #         use_hungarian=use_hungarian,
-    #         length_weight=length_weight,
-    #         return_flow=return_flow,
-    #         progress_notifier=progress_notifier
-    #     )
-
-    #     if return_flow:
-    #         tracking_result, flow_fields = result
-    #     else:
-    #         tracking_result = result
-    #         flow_fields = None
-
-    #     # Map track IDs back to full frame indices
-    #     full_track_ids = [None] * self.metadata.n_stack
-    #     for i, frame_idx in enumerate(list_frames):
-    #         full_track_ids[frame_idx] = tracking_result.track_ids[i]
-
-    #     # Store results in data dict
-    #     tracking_dict = {
-    #         'params.track_sarcomere_vectors.frames': list_frames,
-    #         'params.track_sarcomere_vectors.max_distance': max_distance,
-    #         'params.track_sarcomere_vectors.memory': memory,
-    #         'params.track_sarcomere_vectors.min_track_length': min_track_length,
-    #         'params.track_sarcomere_vectors.flow_blur_sigma': flow_blur_sigma,
-    #         'params.track_sarcomere_vectors.use_hungarian': use_hungarian,
-    #         'params.track_sarcomere_vectors.length_weight': length_weight,
-    #         'track_ids': full_track_ids,
-    #         'n_tracks': tracking_result.n_tracks,
-    #         'track_lengths': tracking_result.track_lengths,
-    #     }
-    #     self.data.update(tracking_dict)
-
-    #     if self.auto_save:
-    #         self.store_structure_data()
-
-    #     logger.info(
-    #         f"Tracking complete: {tracking_result.n_tracks} tracks, "
-    #         f"median length {np.median(tracking_result.track_lengths) if len(tracking_result.track_lengths) > 0 else 0:.0f} frames"
-    #     )
-
-    #     if return_flow:
-    #         return flow_fields
-    #     return None
+    def track_sarcomere_vectors(
+        self,
+        frames: Union[str, int, List[int], np.ndarray] = 'all',
+        max_distance: float = 0.5,
+        memory: int = 5,
+        min_track_length: int = 10,
+        flow_blur_sigma: Optional[float] = None,
+        use_hungarian: bool = True,
+        length_weight: float = 2.0,
+        return_flow: bool = False,
+        progress_notifier: ProgressNotifier = ProgressNotifier.progress_notifier_tqdm()
+    ) -> Optional[List[np.ndarray]]:
+        """
+        Track sarcomere vectors across frames using optical flow.
+
+        This method uses dense optical flow computed on Z-band and M-band masks
+        to predict vector motion between frames, then matches predicted positions
+        to detected vectors using greedy exclusive assignment.
+
+        Parameters
+        ----------
+        frames : {'all', int, list, np.ndarray}, optional
+            Frames for tracking ('all' for all frames, int for a single frame,
+            list or ndarray for selected frames). Defaults to 'all'.
+        max_distance : float, optional
+            Maximum matching distance in µm. Vectors predicted to be further
+            than this from any detection are considered lost. Default 0.5 µm.
+        memory : int, optional
+            Number of frames to keep tracks "sleeping" before termination.
+            Sleeping tracks can be re-linked if a matching vector reappears.
+            Default 5 frames.
+        min_track_length : int, optional
+            Minimum track length to keep. Shorter tracks are filtered out
+            and their vectors marked as untracked (track_id = -1). Default 10.
+        flow_blur_sigma : float, optional
+            Gaussian blur sigma (in pixels) to apply to masks before optical
+            flow computation. Can help with noisy masks. Default None (no blur).
+        use_hungarian : bool, optional
+            If True, use Hungarian algorithm for optimal matching. This gives
+            better results when vectors are densely packed. Default True.
+        length_weight : float, optional
+            Weight for sarcomere length difference in matching cost when using
+            Hungarian algorithm. Higher values favor matching vectors with
+            similar sarcomere lengths. Default 2.0.
+        return_flow : bool, optional
+            If True, also return the computed optical flow fields.
+            Default False.
+        progress_notifier : ProgressNotifier, optional
+            Progress notifier for tracking progress.
+
+        Returns
+        -------
+        flow_fields : list of np.ndarray, optional
+            Only returned if return_flow=True. List of flow fields,
+            each of shape (H, W, 2).
+
+        Notes
+        -----
+        Results are stored in `self.data` with the following keys:
+
+        - ``'track_ids'``: list of (N_t,) arrays with track ID for each vector
+          per frame. Track ID of -1 indicates untracked vectors.
+        - ``'n_tracks'``: total number of unique trajectories.
+        - ``'track_lengths'``: (n_tracks,) array with number of frames each
+          track spans.
+        - ``'params.track_sarcomere_vectors.*'``: tracking parameters used.
+
+        Requires :meth:`analyze_sarcomere_vectors` to be run first.
+
+        Examples
+        --------
+        >>> structure.analyze_sarcomere_vectors(frames='all')
+        >>> structure.track_sarcomere_vectors(max_distance=0.5, memory=5)
+        >>> track_ids = structure.data['track_ids']
+        >>> # track_ids[frame][vector_idx] gives the track ID for that vector
+        """
+        # Validate prerequisites
+        if 'pos_vectors_px' not in self.data:
+            raise ValueError(
+                "Sarcomere vectors not yet analyzed. Run analyze_sarcomere_vectors first."
+            )
+
+        if not os.path.exists(self.file_zbands) or not os.path.exists(self.file_mbands):
+            raise FileNotFoundError(
+                "Z-band or M-band mask not found. Please run detect_sarcomeres first."
+            )
+
+        # Determine frames to track
+        analyzed_frames = self.data.get('params.analyze_sarcomere_vectors.frames', [])
+        if isinstance(frames, str) and frames == 'all':
+            list_frames = list(range(self.metadata.n_stack))
+        elif np.issubdtype(type(frames), np.integer):
+            list_frames = [frames]
+        elif isinstance(frames, (list, np.ndarray)):
+            list_frames = [int(f) for f in frames]
+        else:
+            raise ValueError('frames argument not valid')
+
+        # Validate frames have been analyzed
+        if not set(list_frames).issubset(set(analyzed_frames)):
+            missing = set(list_frames) - set(analyzed_frames)
+            raise ValueError(
+                f"Frames {missing} have not been analyzed. "
+                f"Run analyze_sarcomere_vectors first for these frames."
+            )
+
+        logger.info(f"Starting sarcomere vector tracking for {len(list_frames)} frames...")
+
+        # Load masks
+        zbands_stack = tifffile.imread(self.file_zbands)
+        mbands_stack = tifffile.imread(self.file_mbands)
+
+        # Ensure 3D
+        if zbands_stack.ndim == 2:
+            zbands_stack = np.expand_dims(zbands_stack, axis=0)
+        if mbands_stack.ndim == 2:
+            mbands_stack = np.expand_dims(mbands_stack, axis=0)
+
+        # Select frames
+        zbands_stack = zbands_stack[list_frames]
+        mbands_stack = mbands_stack[list_frames]
+
+        # Prepare vector data for selected frames
+        positions_per_frame = [self.data['pos_vectors_px'][f] for f in list_frames]
+        lengths_per_frame = [self.data['sarcomere_length_vectors'][f] for f in list_frames]
+        orientations_per_frame = [self.data['sarcomere_orientation_vectors'][f] for f in list_frames]
+
+        # Run tracking
+        result = sarcomere_tracking.track_sarcomere_vectors(
+            zbands_stack=zbands_stack,
+            mbands_stack=mbands_stack,
+            positions_per_frame=positions_per_frame,
+            sarcomere_lengths_per_frame=lengths_per_frame,
+            orientations_per_frame=orientations_per_frame,
+            pixelsize=self.metadata.pixelsize,
+            max_distance=max_distance,
+            memory=memory,
+            min_track_length=min_track_length,
+            flow_blur_sigma=flow_blur_sigma,
+            use_hungarian=use_hungarian,
+            length_weight=length_weight,
+            return_flow=return_flow,
+            progress_notifier=progress_notifier
+        )
+
+        if return_flow:
+            tracking_result, flow_fields = result
+        else:
+            tracking_result = result
+            flow_fields = None
+
+        # Map track IDs back to full frame indices
+        full_track_ids = [None] * self.metadata.n_stack
+        for i, frame_idx in enumerate(list_frames):
+            full_track_ids[frame_idx] = tracking_result.track_ids[i]
+
+        # Store results in data dict
+        tracking_dict = {
+            'params.track_sarcomere_vectors.frames': list_frames,
+            'params.track_sarcomere_vectors.max_distance': max_distance,
+            'params.track_sarcomere_vectors.memory': memory,
+            'params.track_sarcomere_vectors.min_track_length': min_track_length,
+            'params.track_sarcomere_vectors.flow_blur_sigma': flow_blur_sigma,
+            'params.track_sarcomere_vectors.use_hungarian': use_hungarian,
+            'params.track_sarcomere_vectors.length_weight': length_weight,
+            'track_ids': full_track_ids,
+            'n_tracks': tracking_result.n_tracks,
+            'track_lengths': tracking_result.track_lengths,
+        }
+        self.data.update(tracking_dict)
+
+        if self.auto_save:
+            self.store_structure_data()
+
+        logger.info(
+            f"Tracking complete: {tracking_result.n_tracks} tracks, "
+            f"median length {np.median(tracking_result.track_lengths) if len(tracking_result.track_lengths) > 0 else 0:.0f} frames"
+        )
+
+        if return_flow:
+            return flow_fields
+        return None
 
     def full_analysis_structure(self, frames='all'):
         """