extract_bboxes_pt21.py make sure to cat along last dimension

Signed-off-by: Martin <[email protected]>
bmmtstb · Aug 5, 2024 · bf56946 · bf56946
1 parent 31c966b
commit bf56946
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/scripts/helpers/extract_bboxes_pt21.py b/scripts/helpers/extract_bboxes_pt21.py
@@ -55,8 +55,8 @@ def save_crops(_s: State, img_dir: FilePath, _gt_img_id: str | int) -> None:
             weights = _s.joint_weight[i].unsqueeze(0).cpu()
         else:
             weights = t.ones((1, _s.J, 1), dtype=t.float32)
-        kp_loc = t.cat([_s.keypoints_local[i].unsqueeze(0).cpu(), weights])
-        kp_glob = t.cat([_s.keypoints[i].unsqueeze(0).cpu(), weights])
+        kp_loc = t.cat([_s.keypoints_local[i].unsqueeze(0).cpu(), weights], dim=-1)
+        kp_glob = t.cat([_s.keypoints[i].unsqueeze(0).cpu(), weights], dim=-1)
         t.save(kp_loc, str(img_path).replace(".jpg", ".pt"))
         t.save(kp_glob, str(img_path).replace(".jpg", "_glob.pt"))