Luigi
commited on
Commit
•
98d66d4
1
Parent(s):
0bf1eb7
Deal with multiple cameras
Browse files- 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 |
-
|
594 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
595 |
for i, (keypoints, scores) in enumerate(zip(b_keypoints, b_scores)):
|
596 |
bboxes = b_bboxes[i]
|
597 |
-
|
598 |
-
self.
|
599 |
|
600 |
frame, bboxes, keypoints, scores = None, None, None, None
|
601 |
-
if not
|
602 |
-
bboxes, keypoints, scores =
|
603 |
-
frame =
|
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.
|
623 |
-
self.
|
624 |
-
self.
|
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
|