Remove (or replace) Instance when Labels.remove_instance

sleap/io/dataset.py

@@ -62,7 +62,7 @@
 import numpy as np
 import datetime
 from sklearn.model_selection import train_test_split
-from sleap.io.cameras import RecordingSession
+from sleap.io.cameras import FrameGroup, RecordingSession
 
 try:
     from typing import ForwardRef
@@ -1405,11 +1405,41 @@ def track_swap(
     def remove_instance(
         self, frame: LabeledFrame, instance: Instance, in_transaction: bool = False
     ):
-        """Remove instance from frame, updating track occupancy."""
+        """Remove instance from frame, updating track occupancy and `FrameGroup`."""
+
         frame.instances.remove(instance)
         if not in_transaction:
             self._cache.remove_instance(frame, instance)
 
+        # Also remove instance from `InstanceGroup` if any
+        session = self.get_session(frame.video)
+        if session is None:
+            return  # No session, so no `FrameGroup` to update
+        frame_group: FrameGroup = session.frame_groups.get(frame.frame_idx, None)
+        if frame_group is None:
+            return  # No `FrameGroup` for this frame
+        video = frame.video
+        camera = session.get_camera(video=video)
+        if camera is None:
+            return  # No camera, so no `InstanceGroup` to update
+
+        # If `Instance.from_predicted`, then replace with `PredictedInstance`
+        instance_group = frame_group.get_instance_group(instance=instance)
+        replace_instance = (instance.from_predicted is not None) and (
+            instance_group is not None
+        )
+
+        # Add the new instance to the `FrameGroup` and replace the old one
+        if replace_instance:
+            frame_group.add_instance(
+                instance=instance.from_predicted,
+                camera=camera,
+                instance_group=instance_group,
+            )
+        # Otherwise just remove the instance
+        else:
+            frame_group.remove_instance(instance=instance)
+
     def add_instance(self, frame: LabeledFrame, instance: Instance):
         """Add instance to frame, updating track occupancy."""
         # Ensure that there isn't already an Instance with this track