Internalize input frame buffer

rtmo_demo.py

@@ -6,7 +6,6 @@ from pathlib import Path
 import argparse
 import os
 from rtmo_gpu import RTMO_GPU_Batch, draw_skeleton, resize_to_fit_screen
-from queue import Queue
 
 if __name__ == "__main__":
 
@@ -28,18 +27,16 @@ if __name__ == "__main__":
     
         # Now, use the best.url, which is the direct video link for streaming
         cap = cv2.VideoCapture(filename=os.path.abspath(mp4_path))
-        in_queue = Queue(maxsize=args.batch_size)
         frame_idx = 0
         s = time.time()
         while cap.isOpened():
             success, frame = cap.read()
             frame_idx += 1
-            in_queue.put(frame)
 
             if not success:
                 break
 
-            keypoints, scores = body(frame)
+            frame_out, keypoints, scores = body(frame)
 
             if keypoints is not None:
                 if frame_idx % args.batch_size == 0 and frame_idx:
@@ -49,8 +46,7 @@ if __name__ == "__main__":
                     print(f'det: {fps} FPS')
                     s = current_time
 
-                frame = in_queue.get()
-                img_show = frame.copy()
+                img_show = frame_out.copy()
                 
                 # if you want to use black background instead of original image,
                 # img_show = np.zeros(img_show.shape, dtype=np.uint8)

rtmo_gpu.py

@@ -579,6 +579,7 @@ class RTMO_GPU_Batch(RTMO_GPU):
     
     def __call__(self, image: np.array):
         self.buffer.append(image)
+        self.in_queue.put(image)
 
         if len(self.buffer) == self.batch_size:
             b_keypoints, b_scores = self.__batch_call__(self.buffer)
@@ -586,11 +587,12 @@ class RTMO_GPU_Batch(RTMO_GPU):
                 self.out_queue.put((keypoints, scores))
             self.buffer = []
 
-        keypoints, scores = None, None
+        frame, keypoints, scores = None, None, None
         if not self.out_queue.empty():
             keypoints, scores = self.out_queue.get()
+            frame = self.in_queue.get()
     
-        return keypoints, scores
+        return frame, keypoints, scores
 
     
     def __init__(self,
@@ -607,6 +609,7 @@ class RTMO_GPU_Batch(RTMO_GPU):
                          is_yolo_nas_pose)
         
         self.batch_size = batch_size
+        self.in_queue = Queue(maxsize=self.batch_size)
         self.out_queue = Queue(maxsize=self.batch_size)
         self.buffer = []