Merge branch 'develop' into track-progress

docs/guides/cli.md

@@ -207,7 +207,7 @@ optional arguments:
   --tracking.clean_iou_threshold TRACKING.CLEAN_IOU_THRESHOLD
                         IOU to use when culling instances *after* tracking. (default: 0)
   --tracking.similarity TRACKING.SIMILARITY
-                        Options: instance, centroid, iou (default: instance)
+                        Options: instance, normalized_instance, object_keypoint, centroid, iou (default: instance)
   --tracking.match TRACKING.MATCH
                         Options: hungarian, greedy (default: greedy)
   --tracking.robust TRACKING.ROBUST

docs/guides/proofreading.md

@@ -50,6 +50,8 @@ There are currently three methods for matching instances in frame N against thes
 - “**centroid**” measures similarity by the distance between the instance centroids
 - “**iou**” measures similarity by the intersection/overlap of the instance bounding boxes
 - “**instance**” measures similarity by looking at the distances between corresponding nodes in the instances, normalized by the number of valid nodes in the candidate instance.
+- “**normalized_instance**” measures similarity by looking at the distances between corresponding nodes in the instances, normalized by the number of valid nodes in the candidate instance and the keypoints normalized by the image size.
+- “**object_keypoint**” measures similarity by measuring the distance between each keypoints from a reference instance and a query instance, takes the exp(-d**2), sum for all the keypoints and divide by the number of visible keypoints in the reference instance.
 
 Once SLEAP has measured the similarity between all the candidates and the instances in frame N, you need to choose a way to pair them up. You can do this either by picking the best match, and the picking the best remaining match for each remaining instance in turn—this is “**greedy**” matching—or you can find the way of matching identities which minimizes the total cost (or: maximizes the total similarity)—this is “**Hungarian**” matching.
 

environment.yml

@@ -10,7 +10,7 @@ channels:
 
 dependencies:
     # Packages SLEAP uses directly
-    - conda-forge::attrs >=21.2.0  #,<=21.4.0
+    - conda-forge::attrs >=21.2.0
     - conda-forge::cattrs ==1.1.1
     - conda-forge::imageio-ffmpeg # Required for imageio to read/write videos with ffmpeg
     - conda-forge::jsmin

environment_no_cuda.yml

@@ -11,7 +11,7 @@ channels:
 
 dependencies:
     # Packages SLEAP uses directly
-    - conda-forge::attrs >=21.2.0  #,<=21.4.0
+    - conda-forge::attrs >=21.2.0
     - conda-forge::cattrs ==1.1.1
     - conda-forge::imageio-ffmpeg # Required for imageio to read/write videos with ffmpeg
     - conda-forge::jsmin

pypi_requirements.txt

@@ -6,6 +6,8 @@
 # These are also distributed through conda and not pip installed when using conda.
 attrs>=21.2.0,<=21.4.0
 cattrs==1.1.1
+imageio
+imageio-ffmpeg
 # certifi>=2017.4.17,<=2021.10.8
 jsmin
 jsonpickle==1.2

sleap/config/pipeline_form.yaml

@@ -435,7 +435,7 @@ inference:
       label: Similarity Method
       type: list
       default: instance
-      options: "instance,centroid,iou,object keypoint"
+      options: "instance,normalized_instance,centroid,iou,object keypoint"
     - name: tracking.match
       label: Matching Method
       type: list
@@ -536,7 +536,7 @@ inference:
       label: Similarity Method
       type: list
       default: instance
-      options: "instance,centroid,iou,object keypoint"
+      options: "instance,normalized_instance,centroid,iou,object keypoint"
     - name: tracking.match
       label: Matching Method
       type: list

sleap/gui/app.py

@@ -377,7 +377,9 @@ def add_menu_item(menu, key: str, name: str, action: Callable):
         def connect_check(key):
             self._menu_actions[key].setCheckable(True)
             self._menu_actions[key].setChecked(self.state[key])
-            self.state.connect(key, self._menu_actions[key].setChecked)
+            self.state.connect(
+                key, lambda checked: self._menu_actions[key].setChecked(checked)
+            )
 
         # add checkable menu item connected to state variable
         def add_menu_check_item(menu, key: str, name: str):

sleap/gui/dataviews.py

@@ -413,13 +413,6 @@ def set_item(self, item, key, value):
         elif key == "symmetry":
             self.context.setNodeSymmetry(skeleton=self.obj, node=item, symmetry=value)
 
-    def get_item_color(self, item: Any, key: str):
-        if self.skeleton:
-            color = self.context.app.color_manager.get_item_color(
-                item, parent_skeleton=self.skeleton
-            )
-            return QtGui.QColor(*color)
-
 
 class SkeletonEdgesTableModel(GenericTableModel):
     """Table model for skeleton edges."""
@@ -436,14 +429,6 @@ def object_to_items(self, skeleton: Skeleton):
             ]
         return items
 
-    def get_item_color(self, item: Any, key: str):
-        if self.skeleton:
-            edge_pair = (item["source"], item["destination"])
-            color = self.context.app.color_manager.get_item_color(
-                edge_pair, parent_skeleton=self.skeleton
-            )
-            return QtGui.QColor(*color)
-
 
 class LabeledFrameTableModel(GenericTableModel):
     """Table model for listing instances in labeled frame.

sleap/gui/widgets/docks.py

@@ -30,10 +30,8 @@
 )
 from sleap.gui.dialogs.formbuilder import YamlFormWidget
 from sleap.gui.widgets.views import CollapsibleWidget
-from sleap.skeleton import Skeleton
-from sleap.util import decode_preview_image, find_files_by_suffix, get_package_file
-
-# from sleap.gui.app import MainWindow
+from sleap.skeleton import Skeleton, SkeletonDecoder
+from sleap.util import find_files_by_suffix, get_package_file
 
 
 class DockWidget(QDockWidget):
@@ -365,7 +363,7 @@ def create_templates_groupbox(self) -> QGroupBox:
         def updatePreviewImage(preview_image_bytes: bytes):
 
             # Decode the preview image
-            preview_image = decode_preview_image(preview_image_bytes)
+            preview_image = SkeletonDecoder.decode_preview_image(preview_image_bytes)
 
             # Create a QImage from the Image
             preview_image = QtGui.QImage(

sleap/nn/inference.py

@@ -2651,6 +2651,7 @@ def _object_builder():
                         # Set tracks for predicted instances in this frame.
                         predicted_instances = self.tracker.track(
                             untracked_instances=predicted_instances,
+                            img_hw=ex["image"].shape[-3:-1],
                             img=image,
                             t=frame_ind,
                         )
@@ -3293,6 +3294,7 @@ def _object_builder():
                         # Set tracks for predicted instances in this frame.
                         predicted_instances = self.tracker.track(
                             untracked_instances=predicted_instances,
+                            img_hw=ex["image"].shape[-3:-1],
                             img=image,
                             t=frame_ind,
                         )

sleap/nn/tracker/components.py

@@ -12,6 +12,7 @@
 
 
 """
+
 import operator
 from collections import defaultdict
 import logging
@@ -29,6 +30,21 @@
 InstanceType = TypeVar("InstanceType", Instance, PredictedInstance)
 
 
+def normalized_instance_similarity(
+    ref_instance: InstanceType, query_instance: InstanceType, img_hw: Tuple[int]
+) -> float:
+    """Computes similarity between instances with normalized keypoints."""
+
+    normalize_factors = np.array((img_hw[1], img_hw[0]))
+    ref_visible = ~(np.isnan(ref_instance.points_array).any(axis=1))
+    normalized_query_keypoints = query_instance.points_array / normalize_factors
+    normalized_ref_keypoints = ref_instance.points_array / normalize_factors
+    dists = np.sum((normalized_query_keypoints - normalized_ref_keypoints) ** 2, axis=1)
+    similarity = np.nansum(np.exp(-dists)) / np.sum(ref_visible)
+
+    return similarity
+
+
 def instance_similarity(
     ref_instance: InstanceType, query_instance: InstanceType
 ) -> float:

sleap/nn/tracking.py

@@ -11,12 +11,14 @@
 import cv2
 import numpy as np
 import rich.progress
+import functools
 
 from sleap import Track, LabeledFrame, Skeleton
 
 from sleap.nn.tracker.components import (
     factory_object_keypoint_similarity,
     instance_similarity,
+    normalized_instance_similarity,
     centroid_distance,
     instance_iou,
     hungarian_matching,
@@ -504,7 +506,8 @@ def get_candidates(
     instance=instance_similarity,
     centroid=centroid_distance,
     iou=instance_iou,
-    object_keypoint=instance_similarity,
+    normalized_instance=normalized_instance_similarity,
+    object_keypoint=factory_object_keypoint_similarity,
 )
 
 match_policies = dict(
@@ -799,19 +802,26 @@ def uses_image(self):
     def track(
         self,
         untracked_instances: List[InstanceType],
+        img_hw: Tuple[int],
         img: Optional[np.ndarray] = None,
         t: int = None,
     ) -> List[InstanceType]:
         """Performs a single step of tracking.
 
         Args:
             untracked_instances: List of instances to assign to tracks.
+            img_hw: (height, width) of the image used to normalize the keypoints.
             img: Image data of the current frame for flow shifting.
             t: Current timestep. If not provided, increments from the internal queue.
 
         Returns:
             A list of the instances that were tracked.
         """
+        if self.similarity_function == normalized_instance_similarity:
+            factory_normalized_instance = functools.partial(
+                normalized_instance_similarity, img_hw=img_hw
+            )
+            self.similarity_function = factory_normalized_instance
 
         if self.candidate_maker is None:
             return untracked_instances