Merge branch 'develop' into elizabeth/update-python-and-dependencies

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.
 

sleap/config/pipeline_form.yaml

@@ -439,7 +439,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
@@ -538,7 +538,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/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

@@ -2622,6 +2622,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,
                         )
@@ -3264,6 +3265,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

@@ -5,13 +5,15 @@
 import attr
 import numpy as np
 import cv2
+import functools
 from typing import Callable, Deque, Dict, Iterable, List, Optional, Tuple
 
 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,
@@ -495,7 +497,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(
@@ -639,19 +642,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
@@ -1520,6 +1530,7 @@ def run_tracker(frames: List[LabeledFrame], tracker: BaseTracker) -> List[Labele
             track_args["img"] = lf.video[lf.frame_idx]
         else:
             track_args["img"] = None
+        track_args["img_hw"] = lf.image.shape[-3:-1]
 
         new_lf = LabeledFrame(
             frame_idx=lf.frame_idx,