fix errors from merge

talmolab · May 7, 2024 · 9edc517 · 9edc517
1 parent 390a564
commit 9edc517
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 40 deletions.
diff --git a/biogtr/data_structures.py b/biogtr/data_structures.py
@@ -635,7 +635,7 @@ def to(self, map_location: str):
  return self
 
  def to_slp(
- self, track_lookup: dict[int : sio.Track] = {}
+ self, track_lookup: dict[int, sio.Track] = {}
  ) -> tuple[sio.LabeledFrame, dict[int, sio.Track]]:
  """Convert Frame to sleap_io.LabeledFrame object.
 

diff --git a/biogtr/datasets/sleap_dataset.py b/biogtr/datasets/sleap_dataset.py
@@ -24,7 +24,6 @@ def __init__(
  padding: int = 5,
  crop_size: int = 128,
  anchors: Union[int, list[str], str] = "",
- anchors: Union[int, list[str], str] = "",
  chunk: bool = True,
  clip_length: int = 500,
  mode: str = "train",
@@ -99,7 +98,6 @@ def __init__(
  ) or self.anchors == 0:
  raise ValueError(f"Must provide at least one anchor but got {self.anchors}")
 
-
  if isinstance(anchors, int):
  self.anchors = anchors
  elif isinstance(anchors, str):
@@ -296,11 +294,6 @@ def get_instances(self, label_idx: List[int], frame_idx: List[int]) -> list[dict
  else:
  anchors = self.anchors
 
- for anchor in anchors:
- if anchor == "midpoint" or anchor == "centroid":
- centroid = np.nanmean(np.array(list(pose.values())), axis=0)
- anchors = self.anchors
-
  for anchor in anchors:
  if anchor == "midpoint" or anchor == "centroid":
  centroid = np.nanmean(np.array(list(pose.values())), axis=0)
@@ -313,6 +306,7 @@ def get_instances(self, label_idx: List[int], frame_idx: List[int]) -> list[dict
  elif anchor not in pose and len(anchors) == 1:
  anchor = "midpoint"
  centroid = np.nanmean(np.array(list(pose.values())), axis=0)
+
  elif anchor in pose:
  centroid = np.array(pose[anchor])
  if np.isnan(centroid).any():
@@ -325,18 +319,9 @@ def get_instances(self, label_idx: List[int], frame_idx: List[int]) -> list[dict
  else:
  centroid = np.array([np.nan, np.nan])
 
- else:
- centroid = np.array([np.nan, np.nan])
-
- if np.isnan(centroid).all():
- bbox = torch.tensor([np.nan, np.nan, np.nan, np.nan])
- else:
- bbox = data_utils.pad_bbox(
- data_utils.get_bbox(centroid, self.crop_size),
- padding=self.padding,
- )
  if np.isnan(centroid).all():
  bbox = torch.tensor([np.nan, np.nan, np.nan, np.nan])
+
  else:
  bbox = data_utils.pad_bbox(
  data_utils.get_bbox(centroid, self.crop_size),
@@ -352,25 +337,6 @@ def get_instances(self, label_idx: List[int], frame_idx: List[int]) -> list[dict
  )
  else:
  crop = data_utils.crop_bbox(img, bbox)
- if bbox.isnan().all():
- crop = torch.zeros(
- c,
- self.crop_size + 2 * self.padding,
- self.crop_size + 2 * self.padding,
- dtype=img.dtype,
- )
- else:
- crop = data_utils.crop_bbox(img, bbox)
-
- crops.append(crop)
- centroids[anchor] = centroid
- boxes.append(bbox)
-
- if len(crops) > 0:
- crops = torch.concat(crops, dim=0)
-
- if len(boxes) > 0:
- boxes = torch.stack(boxes, dim=0)
 
  crops.append(crop)
  centroids[anchor] = centroid
@@ -388,9 +354,6 @@ def get_instances(self, label_idx: List[int], frame_idx: List[int]) -> list[dict
  crop=crops,
  centroid=centroids,
  bbox=boxes,
- crop=crops,
- centroid=centroids,
- bbox=boxes,
  skeleton=skeleton,
  pose=poses[j],
  point_scores=point_scores[j],