Fix instance color (#2434)

* Fix instance color

* Fix cli cmd

* Minor bug fix

Author: gitttt-1234

sleap/gui/color.py

@@ -19,7 +19,7 @@
 from sleap.util import get_config_file
 from sleap_io.model.instance import Instance, Track
 from sleap_io.model.skeleton import Node, Edge
-from sleap_io import Labels
+from sleap_io import Labels, LabeledFrame
 from sleap_io.model.skeleton import Skeleton
 from sleap.prefs import prefs
 from sleap.sleap_io_adaptors.skeleton_utils import node_to_index, edge_to_index
@@ -248,6 +248,7 @@ def get_item_color(
         item: Any,
         parent_instance: Optional[Instance] = None,
         parent_skeleton: Optional[Skeleton] = None,
+        frame: Optional[LabeledFrame] = None,
     ) -> ColorTupleType:
         """Gets (r, g, b) tuple of color to use for drawing item."""
 
@@ -276,7 +277,7 @@ def get_item_color(
 
             if track is None and parent_instance:
                 # Get an index for items without track
-                track = self.get_pseudo_track_index(parent_instance)
+                track = self.get_pseudo_track_index(parent_instance, frame=frame)
 
             return self.get_track_color(track=track)
 

sleap/gui/learning/runners.py

@@ -500,11 +500,13 @@ def write_pipeline_files(
                 # Add a line to the script for training this model
                 train_script += (
                     f"sleap-nn-train --config-name {new_cfg_filename} "
-                    f"--config-dir {''} "
+                    f"--config-dir . "
                     f"trainer_config.ckpt_dir={Path(ckpt_path).parent.as_posix()} "
-                    f"trainer_config.run_name={Path(ckpt_path).name}"
-                    f"trainer_config.zmq.controller_port={cfg_info.config.trainer_config.zmq.controller_port}"
-                    f"trainer_config.zmq.publish_port={cfg_info.config.trainer_config.zmq.publish_port}"
+                    f"trainer_config.run_name={Path(ckpt_path).name} "
+                    f"trainer_config.zmq.controller_port="
+                    f"{cfg_info.config.trainer_config.zmq.controller_port} "
+                    f"trainer_config.zmq.publish_port="
+                    f"{cfg_info.config.trainer_config.zmq.publish_port} "
                     "\n"
                 )
 

sleap/gui/overlays/instance.py

@@ -45,6 +45,7 @@ def add_to_scene(self, video, frame_idx):
         for instance in instances:
             self.player.addInstance(
                 instance=instance,
+                frame=lf,
                 markerRadius=self.state.get("marker size", 4),
                 nodeLabelSize=self.state.get("node label size", 12),
                 show_non_visible=self.state.get("show non-visible nodes", default=True),

sleap/gui/widgets/video.py

@@ -60,7 +60,7 @@
 from sleap_io.model.instance import Instance, PredictedInstance
 from sleap.sleap_io_adaptors.instance_utils import fill_missing, node_points
 from sleap.sleap_io_adaptors.video_utils import get_last_frame_idx
-from sleap_io import Video
+from sleap_io import Video, LabeledFrame
 from sleap.prefs import prefs
 from sleap_io import Node
 
@@ -484,7 +484,7 @@ def scene(self):
         """Returns `QGraphicsScene` for viewer."""
         return self.view.scene
 
-    def addInstance(self, instance, **kwargs):
+    def addInstance(self, instance, frame: Optional[LabeledFrame] = None, **kwargs):
         """Add a skeleton instance to the video.
 
         Args:
@@ -494,7 +494,7 @@ def addInstance(self, instance, **kwargs):
         """
         # Check if instance is an Instance (or subclass of Instance)
         if issubclass(type(instance), Instance):
-            instance = QtInstance(instance=instance, player=self, **kwargs)
+            instance = QtInstance(instance=instance, frame=frame, player=self, **kwargs)
         if type(instance) != QtInstance:
             return
         if instance.instance.n_visible > 0 or not isinstance(
@@ -1415,6 +1415,7 @@ def __init__(
         predicted=False,
         show_non_visible=True,
         callbacks=None,
+        frame: Optional[LabeledFrame] = None,
         *args,
         **kwargs,
     ):
@@ -1425,7 +1426,7 @@ def __init__(
         self.radius = radius
         self.color_manager = self.player.color_manager
         self.color = self.color_manager.get_item_color(
-            self.node, self._parent_instance.instance
+            self.node, self._parent_instance.instance, frame=frame
         )
         self.edges = []
         self.name = node.name
@@ -1676,6 +1677,7 @@ def __init__(
         src: QtNode,
         dst: QtNode,
         show_non_visible: bool = True,
+        frame: Optional[LabeledFrame] = None,
         *args,
         **kwargs,
     ):
@@ -1702,7 +1704,9 @@ def __init__(
         )
 
         edge_pair = (src.node, dst.node)
-        color = player.color_manager.get_item_color(edge_pair, parent.instance)
+        color = player.color_manager.get_item_color(
+            edge_pair, parent.instance, frame=frame
+        )
         pen_width = player.color_manager.get_item_pen_width(edge_pair, parent.instance)
         pen = QPen(QColor(*color), pen_width)
         pen.setCosmetic(True)
@@ -1834,6 +1838,7 @@ def __init__(
         markerRadius=4,
         nodeLabelSize=12,
         show_non_visible=True,
+        frame: Optional[LabeledFrame] = None,
         *args,
         **kwargs,
     ):
@@ -1844,7 +1849,7 @@ def __init__(
         self.predicted = hasattr(instance, "score")
 
         color_manager = self.player.color_manager
-        color = color_manager.get_item_color(self.instance)
+        color = color_manager.get_item_color(self.instance, frame=frame)
 
         self.show_non_visible = show_non_visible
         self.selectable = not self.predicted or color_manager.color_predicted
@@ -1875,7 +1880,9 @@ def __init__(
         if self.predicted:
             self.box = QGraphicsRectItem(parent=self)
         else:
-            self.box = VisibleBoundingBox(rect=self._bounding_rect, parent=self)
+            self.box = VisibleBoundingBox(
+                rect=self._bounding_rect, parent=self, frame=frame
+            )
         box_pen_width = color_manager.get_item_pen_width(self.instance)
         box_pen = QPen(QColor(*color), box_pen_width)
         box_pen.setStyle(Qt.DashLine)
@@ -1924,6 +1931,7 @@ def __init__(
                     predicted=self.predicted,
                     radius=self.markerRadius,
                     show_non_visible=self.show_non_visible,
+                    frame=frame,
                 )
 
                 self.nodes[node.name] = node_item
@@ -1938,6 +1946,7 @@ def __init__(
                     src=self.nodes[src],
                     dst=self.nodes[dst],
                     show_non_visible=self.show_non_visible,
+                    frame=frame,
                 )
                 self.nodes[src].edges.append(edge_item)
                 self.nodes[dst].edges.append(edge_item)
@@ -2254,11 +2263,12 @@ def __init__(
         parent: QtInstance,
         opacity: float = 0.8,
         scaling_padding: float = 10.0,
+        frame: Optional[LabeledFrame] = None,
     ):
         super().__init__(rect, parent)
         self.box_width = parent.markerRadius
         color_manager = parent.player.color_manager
-        int_color = color_manager.get_item_color(parent.instance)
+        int_color = color_manager.get_item_color(parent.instance, frame=frame)
         self.int_color = QColor(*int_color)
         self.corner_opacity = opacity
         self.scaling_padding = scaling_padding
@@ -2512,6 +2522,7 @@ def plot_instances(scene, frame_idx, labels, video=None, fixed=True):
             color=color_manager.get_track_color(pseudo_track),
             predicted=fixed,
             color_predicted=True,
+            frame=labeled_frame,
             show_non_visible=False,
         )
         inst.showLabels(False)

sleap/io/visuals.py

@@ -17,7 +17,7 @@
 
 from sleap.gui.color import ColorManager
 from sleap_io.model.instance import Instance
-from sleap_io import Video, Labels
+from sleap_io import Video, Labels, LabeledFrame
 from sleap.sleap_io_adaptors.video_utils import _sentinel
 from sleap.sleap_io_adaptors.lf_labels_utils import (
     load_labels_video_search,
@@ -204,7 +204,7 @@ def _plot_instances_cv(
             img, offset = self._crop_frame(img, instances)
 
         for instance in instances:
-            self._plot_instance_cv(img, instance, offset)
+            self._plot_instance_cv(img, instance, offset, frame=lfs[0])
 
         return img
 
@@ -270,6 +270,7 @@ def _plot_instance_cv(
         instance: "Instance",
         offset: Optional[Tuple[int, int]] = None,
         fill: bool = True,
+        frame: Optional[LabeledFrame] = None,
     ):
         """
         Add visual annotations for single instance.
@@ -299,7 +300,9 @@ def _plot_instance_cv(
 
         for node_idx, (x, y) in enumerate(points_array):
             node = nodes[node_idx]
-            node_color_bgr = self.color_manager.get_item_color(node, instance)[::-1]
+            node_color_bgr = self.color_manager.get_item_color(
+                node, instance, frame=frame
+            )[::-1]
 
             # Make sure this is a valid and visible point
             if not has_nans(x, y):
@@ -323,7 +326,9 @@ def _plot_instance_cv(
                 dst_x, dst_y = points_array[dst]
 
                 edge = (nodes[src], nodes[dst])
-                edge_color_bgr = self.color_manager.get_item_color(edge, instance)[::-1]
+                edge_color_bgr = self.color_manager.get_item_color(
+                    edge, instance, frame=frame
+                )[::-1]
 
                 # Make sure that both nodes are present in this instance before
                 # drawing edge