Add method to free buffers

rtmo_gpu.py

@@ -1,11 +1,14 @@
 import os
 import numpy as np
-from typing import List, Tuple
+from typing import List, Tuple, Union
 import onnxruntime as ort
 import cv2
 from queue import Queue
-os.environ['ORT_TENSORRT_EXTRA_PLUGIN_LIB_PATHS']='libmmdeploy_tensorrt_ops.so'
+
+PLUGIN_LIB_PATHS='libmmdeploy_tensorrt_ops.so'
+os.environ['ORT_TENSORRT_EXTRA_PLUGIN_LIB_PATHS']=PLUGIN_LIB_PATHS
 TRT_BACKEND='POLYGRAPHY'
+DEBUG=False
 
 # dictionary from https://github.com/Tau-J/rtmlib/blob/4b29101d54b611048ef165277cebfffff3030074/rtmlib/visualization/skeleton/coco17.py
 coco17 = dict(name='coco17',
@@ -307,7 +310,7 @@ def get_model_format_and_input_shape(model):
     elif is_trt_engine(model):
         model_format = 'engine'
         from polygraphy.backend.trt import load_plugins
-        load_plugins(plugins=['libmmdeploy_tensorrt_ops.so'])
+        load_plugins(plugins=[PLUGIN_LIB_PATHS])
         input_shape = get_trt_input_shapes(model)['input']
     else:
         raise TypeError("Your model is neither ONNX nor Engine !")
@@ -317,10 +320,8 @@ class RTMO_GPU(object):
 
     def preprocess(self, img: np.ndarray):
         """Do preprocessing for RTMPose model inference.
-
         Args:
             img (np.ndarray): Input image in shape.
-
         Returns:
             tuple:
             - resized_img (np.ndarray): Preprocessed image.
@@ -358,17 +359,13 @@ class RTMO_GPU(object):
         ratio: float = 1.,
     ) -> Tuple[np.ndarray, np.ndarray]:
         """Do postprocessing for RTMO model inference.
-
         Args:
             outputs (List[np.ndarray]): Outputs of RTMO model.
             ratio (float): Ratio of preprocessing.
-
         Returns:
             tuple:
             - final_boxes (np.ndarray): Final bounding boxes.
             - final_scores (np.ndarray): Final scores.
-            - final keypoints
-            - final keypoints scores
         """
         
         if not self.is_yolo_nas_pose:
@@ -408,10 +405,8 @@ class RTMO_GPU(object):
 
     def inference(self, img: np.ndarray):
             """Inference model.
-
             Args:
                 img (np.ndarray): Input image in shape.
-
             Returns:
                 outputs (np.ndarray): Output of RTMPose model.
             """
@@ -496,7 +491,7 @@ class RTMO_GPU(object):
                  device: str = 'cuda',
                  is_yolo_nas_pose = False, 
                  batch_size = 1,
-                 plugin_path = 'libmmdeploy_tensorrt_ops.so'):
+                 plugin_path = PLUGIN_LIB_PATHS):
 
         self.batch_size = batch_size
 
@@ -595,10 +590,8 @@ class RTMO_GPU(object):
 class RTMO_GPU_Batch(RTMO_GPU):
     def preprocess_batch(self, imgs: List[np.ndarray]) -> Tuple[np.ndarray, List[float]]:
         """Process a batch of images for RTMPose model inference.
-
         Args:
             imgs (List[np.ndarray]): List of input images.
-
         Returns:
             tuple:
             - batch_img (np.ndarray): Batch of preprocessed images.
@@ -619,10 +612,8 @@ class RTMO_GPU_Batch(RTMO_GPU):
 
     def inference(self, batch_img: np.ndarray):
         """Override to handle batch inference.
-
         Args:
             batch_img (np.ndarray): Batch of preprocessed images.
-
         Returns:
             outputs (List[np.ndarray]): Outputs of RTMPose model for each image.
         """
@@ -690,11 +681,9 @@ class RTMO_GPU_Batch(RTMO_GPU):
         ratios: List[float]
     ) -> Tuple[List[np.ndarray], List[np.ndarray]]:
         """Process outputs for a batch of images.
-
         Args:
             outputs (List[np.ndarray]): Outputs from the model for each image.
             ratios (List[float]): Ratios used for preprocessing each image.
-
         Returns:
             List[Tuple[np.ndarray, np.ndarray]]: keypoints and scores for each image.
         """
@@ -720,6 +709,15 @@ class RTMO_GPU_Batch(RTMO_GPU):
         bboxes, bboxes_scores, keypoints, scores = self.postprocess_batch(outputs, ratios)
         return bboxes, bboxes_scores, keypoints, scores
     
+    def free_unused_buffers(self, activate_cameras_ids: List):
+        for camera_id in list(self.buffers.keys()):
+            if camera_id not in activate_cameras_ids:
+                del self.buffers[camera_id]
+                del self.in_queues[camera_id]
+                del self.out_queues[camera_id]
+                if DEBUG:
+                    print(f'RTMO buffers to camera "{camera_id}" got freed.', flush=True)
+
     def __call__(self, image: np.array, camera_id = 0):
 
         # initialize dedicated buffers & queues for camera with id "camera_id"
@@ -727,6 +725,9 @@ class RTMO_GPU_Batch(RTMO_GPU):
             self.buffers[camera_id] = []
             self.in_queues[camera_id] = Queue(maxsize=self.batch_size)
             self.out_queues[camera_id] = Queue(maxsize=self.batch_size)
+            if DEBUG:
+                print(f'RTMO buffers to camera "{camera_id}" are created.', flush=True)
+
 
         in_queue = self.in_queues[camera_id]
         out_queue = self.out_queues[camera_id]
@@ -755,7 +756,7 @@ class RTMO_GPU_Batch(RTMO_GPU):
                  std: tuple = None,
                  device: str = 'cuda',
                  is_yolo_nas_pose = False,
-                 plugin_path = 'libmmdeploy_tensorrt_ops.so',
+                 plugin_path = PLUGIN_LIB_PATHS,
                  batch_size: int = 1):
         super().__init__(model, 
                          mean, 
@@ -786,4 +787,4 @@ def resize_to_fit_screen(image, screen_width, screen_height):
     # Resize the image
     resized_image = cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_AREA)
 
-    return resized_image
+    return resized_image