Deal with multiple cameras

rtmo_gpu.py

@@ -586,21 +586,30 @@ class RTMO_GPU_Batch(RTMO_GPU):
         bboxes, keypoints, scores = self.postprocess_batch(outputs, ratios)
         return bboxes, keypoints, scores
     
-    def __call__(self, image: np.array):
-        self.buffer.append(image)
-        self.in_queue.put(image)
+    def __call__(self, image: np.array, camera_id = 0):
 
-        if len(self.buffer) == self.batch_size:
-            b_bboxes, b_keypoints, b_scores = self.__batch_call__(self.buffer)
+        # initialize dedicated buffers & queues for camera with id "camera_id"
+        if camera_id not in self.buffers:
+            self.buffers[camera_id] = []
+            self.in_queues[camera_id] = Queue(maxsize=self.batch_size)
+            self.out_queues[camera_id] = Queue(maxsize=self.batch_size)
+
+        in_queue = self.in_queues[camera_id]
+        out_queue = self.out_queues[camera_id]
+        self.buffers[camera_id].append(image)
+        in_queue.put(image)
+
+        if len(self.buffers[camera_id]) == self.batch_size:
+            b_bboxes, b_keypoints, b_scores = self.__batch_call__(self.buffers[camera_id])
             for i, (keypoints, scores) in enumerate(zip(b_keypoints, b_scores)):
                 bboxes = b_bboxes[i]
-                self.out_queue.put((bboxes, keypoints, scores))
-            self.buffer = []
+                out_queue.put((bboxes, keypoints, scores))
+            self.buffers[camera_id] = []
 
         frame, bboxes, keypoints, scores = None, None, None, None
-        if not self.out_queue.empty():
-            bboxes, keypoints, scores = self.out_queue.get()
-            frame = self.in_queue.get()
+        if not out_queue.empty():
+            bboxes, keypoints, scores = out_queue.get()
+            frame = in_queue.get()
     
         return frame, bboxes, keypoints, scores
 
@@ -619,9 +628,9 @@ 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 = []
+        self.in_queues = dict()
+        self.out_queues = dict()
+        self.buffers = dict()
 
 def resize_to_fit_screen(image, screen_width, screen_height):
     # Get the dimensions of the image