Luigi commited on
Commit
6008f96
·
1 Parent(s): 58a44cf

fix postprocess_batch method in RTMO_GPU_Batch class

Browse files
Files changed (1) hide show
  1. rtmo_gpu.py +3 -1
rtmo_gpu.py CHANGED
@@ -547,8 +547,10 @@ class RTMO_GPU_Batch(RTMO_GPU):
547
  batch_keypoints = []
548
  batch_scores = []
549
 
 
550
  for i, ratio in enumerate(ratios):
551
- keypoints, scores = super().postprocess(outputs, ratio)
 
552
  batch_keypoints.append(keypoints)
553
  batch_scores.append(scores)
554
 
 
547
  batch_keypoints = []
548
  batch_scores = []
549
 
550
+ b_dets, b_keypoints = outputs
551
  for i, ratio in enumerate(ratios):
552
+ output = [np.expand_dims(b_dets[i], axis=0), np.expand_dims(b_keypoints[i],axis=0)]
553
+ keypoints, scores = super().postprocess(output, ratio)
554
  batch_keypoints.append(keypoints)
555
  batch_scores.append(scores)
556