fix small bug

dreem/io/frame.py

@@ -91,7 +91,7 @@ def __repr__(self) -> str:
  ")"
  )
 
- def to(self, map_location: str) -> "Frame":
+ def to(self, map_location: Union[str, torch.device]) -> "Frame":
  """Move frame to different device or dtype (See `torch.to` for more info).
 
  Args:
@@ -116,7 +116,7 @@ def to(self, map_location: str) -> "Frame":
  for instance in self.instances:
  instance = instance.to(map_location)
 
- if isinstance(map_location, str):
+ if isinstance(map_location, (str, torch.device)):
  self._device = map_location
 
  return self

dreem/io/instance.py

@@ -146,7 +146,7 @@ def to(self, map_location: Union[str, torch.device]) -> "Instance":
  self._bbox = self._bbox.to(map_location)
  self._crop = self._crop.to(map_location)
  self._features = self._features.to(map_location)
- if isinstance(map_location, str, torch.device):
+ if isinstance(map_location, (str, torch.device)):
  self.device = map_location
 
  return self

dreem/io/visualize.py

@@ -316,12 +316,15 @@ def main(cfg: DictConfig):
  """Take in a path to a video + labels file, annotates a video and saves it to the specified path."""
  labels = pd.read_csv(cfg.labels_path)
  video = imageio.get_reader(cfg.vid_path, "ffmpeg")
- frames_annotated = annotate_video(video, labels, save_path=cfg.save_path, **cfg.annotate)
-
+ frames_annotated = annotate_video(
+ video, labels, save_path=cfg.save_path, **cfg.annotate
+ )
+
  if frames_annotated:
  print("Video saved to {cfg.save_path}!")
  else:
  print("Failed to annotate video!")
 
+
 if __name__ == "__main__":
  main()