Bugfix for batch demo: add padding if remaning number of frames is not equal to batch size
Browse files- demo.sh +1 -1
- demo_batch.sh +1 -1
- rtmo_demo_batch.py +8 -0
demo.sh
CHANGED
@@ -1,2 +1,2 @@
|
|
1 |
#!/bin/sh
|
2 |
-
python3 rtmo_demo.py ./video rtmo-
|
|
|
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-
|
|
|
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]
|