Luigi commited on
Commit
98d66d4
1 Parent(s): 0bf1eb7

Deal with multiple cameras

Browse files
Files changed (1) hide show
  1. rtmo_gpu.py +22 -13
rtmo_gpu.py CHANGED
@@ -586,21 +586,30 @@ class RTMO_GPU_Batch(RTMO_GPU):
586
  bboxes, keypoints, scores = self.postprocess_batch(outputs, ratios)
587
  return bboxes, keypoints, scores
588
 
589
- def __call__(self, image: np.array):
590
- self.buffer.append(image)
591
- self.in_queue.put(image)
592
 
593
- if len(self.buffer) == self.batch_size:
594
- b_bboxes, b_keypoints, b_scores = self.__batch_call__(self.buffer)
 
 
 
 
 
 
 
 
 
 
 
595
  for i, (keypoints, scores) in enumerate(zip(b_keypoints, b_scores)):
596
  bboxes = b_bboxes[i]
597
- self.out_queue.put((bboxes, keypoints, scores))
598
- self.buffer = []
599
 
600
  frame, bboxes, keypoints, scores = None, None, None, None
601
- if not self.out_queue.empty():
602
- bboxes, keypoints, scores = self.out_queue.get()
603
- frame = self.in_queue.get()
604
 
605
  return frame, bboxes, keypoints, scores
606
 
@@ -619,9 +628,9 @@ class RTMO_GPU_Batch(RTMO_GPU):
619
  is_yolo_nas_pose)
620
 
621
  self.batch_size = batch_size
622
- self.in_queue = Queue(maxsize=self.batch_size)
623
- self.out_queue = Queue(maxsize=self.batch_size)
624
- self.buffer = []
625
 
626
  def resize_to_fit_screen(image, screen_width, screen_height):
627
  # Get the dimensions of the image
 
586
  bboxes, keypoints, scores = self.postprocess_batch(outputs, ratios)
587
  return bboxes, keypoints, scores
588
 
589
+ def __call__(self, image: np.array, camera_id = 0):
 
 
590
 
591
+ # initialize dedicated buffers & queues for camera with id "camera_id"
592
+ if camera_id not in self.buffers:
593
+ self.buffers[camera_id] = []
594
+ self.in_queues[camera_id] = Queue(maxsize=self.batch_size)
595
+ self.out_queues[camera_id] = Queue(maxsize=self.batch_size)
596
+
597
+ in_queue = self.in_queues[camera_id]
598
+ out_queue = self.out_queues[camera_id]
599
+ self.buffers[camera_id].append(image)
600
+ in_queue.put(image)
601
+
602
+ if len(self.buffers[camera_id]) == self.batch_size:
603
+ b_bboxes, b_keypoints, b_scores = self.__batch_call__(self.buffers[camera_id])
604
  for i, (keypoints, scores) in enumerate(zip(b_keypoints, b_scores)):
605
  bboxes = b_bboxes[i]
606
+ out_queue.put((bboxes, keypoints, scores))
607
+ self.buffers[camera_id] = []
608
 
609
  frame, bboxes, keypoints, scores = None, None, None, None
610
+ if not out_queue.empty():
611
+ bboxes, keypoints, scores = out_queue.get()
612
+ frame = in_queue.get()
613
 
614
  return frame, bboxes, keypoints, scores
615
 
 
628
  is_yolo_nas_pose)
629
 
630
  self.batch_size = batch_size
631
+ self.in_queues = dict()
632
+ self.out_queues = dict()
633
+ self.buffers = dict()
634
 
635
  def resize_to_fit_screen(image, screen_width, screen_height):
636
  # Get the dimensions of the image