pesi
/

Luigi commited on
Commit
12e168d
1 Parent(s): b75f05d

Bugfix for batch demo: add padding if remaning number of frames is not equal to batch size

Browse files
Files changed (3) hide show
  1. demo.sh +1 -1
  2. demo_batch.sh +1 -1
  3. rtmo_demo_batch.py +8 -0
demo.sh CHANGED
@@ -1,2 +1,2 @@
1
  #!/bin/sh
2
- python3 rtmo_demo.py ./video rtmo-t.fp16_batch1.onnx
 
1
  #!/bin/sh
2
+ python3 rtmo_demo.py ./video rtmo-s.fp16_batch1.onnx
demo_batch.sh CHANGED
@@ -1,2 +1,2 @@
1
  #!/bin/sh
2
- python3 rtmo_demo_batch.py ./video rtmo-t.fp16_batch4.onnx 4
 
1
  #!/bin/sh
2
+ python3 rtmo_demo_batch.py ./video rtmo-s.fp16_batch4.onnx 4
rtmo_demo_batch.py CHANGED
@@ -42,6 +42,14 @@ def process_video(video_path, body_estimator, batch_size=4):
42
 
43
  # Process remaining frames if any
44
  if batch_frames:
 
 
 
 
 
 
 
 
45
  batch_keypoints, batch_scores = body_estimator(batch_frames)
46
  for i, keypoints in enumerate(batch_keypoints):
47
  scores = batch_scores[i]
 
42
 
43
  # Process remaining frames if any
44
  if batch_frames:
45
+ # Padding
46
+ while len(batch_frames) < batch_size:
47
+ # Option 1: Add a black frame
48
+ # black_frame = np.zeros_like(batch_frames[0])
49
+ # batch_frames.append(black_frame)
50
+
51
+ # Option 2: Duplicate the last frame
52
+ batch_frames.append(batch_frames[-1])
53
  batch_keypoints, batch_scores = body_estimator(batch_frames)
54
  for i, keypoints in enumerate(batch_keypoints):
55
  scores = batch_scores[i]