add changes from @coderabbitai's code review

talmolab · Apr 22, 2024 · fc9c16e · fc9c16e
1 parent ec6d359
commit fc9c16e
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 4 deletions.
diff --git a/biogtr/datasets/sleap_dataset.py b/biogtr/datasets/sleap_dataset.py
@@ -6,6 +6,7 @@
 import numpy as np
 import sleap_io as sio
 import random
+import warnings
 from biogtr.data_structures import Frame, Instance
 from biogtr.datasets import data_utils
 from biogtr.datasets.base_dataset import BaseDataset

diff --git a/biogtr/inference/track_queue.py b/biogtr/inference/track_queue.py
@@ -171,8 +171,8 @@ def end_tracks(self, track_id=None):
             try:
                 self._queues.pop(track_id)
                 self._curr_gap.pop(track_id)
-            except Exception as e:
-                print(f"Unable to end track due to {e}")
+            except KeyError:
+                print(f"Track ID {track_id} not found in queue!")
                 return False
         return True
 

diff --git a/biogtr/models/model_utils.py b/biogtr/models/model_utils.py
@@ -1,6 +1,5 @@
 """Module containing model helper functions."""
 
-from copy import deepcopy
 from typing import List, Tuple, Iterable
 from pytorch_lightning import loggers
 from biogtr.data_structures import Frame
@@ -22,7 +21,7 @@ def get_boxes_times(frames: List[Frame]) -> Tuple[torch.Tensor, torch.Tensor]:
     _, h, w = frames[0].img_shape.flatten()
 
     for fidx, frame in enumerate(frames):
-        bbox = deepcopy(frame.get_bboxes())
+        bbox = frame.get_bboxes().clone()
         bbox[:, [0, 2]] /= w
         bbox[:, [1, 3]] /= h