Add new class 'RTMO_GPU_Batch' that can perform inference on batch of images

rtmo_gpu.py

@@ -378,3 +378,87 @@ class RTMO_GPU(object):
         self.std = std
         self.device = device
 
+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.
+            - ratios (List[float]): Ratios used for preprocessing each image.
+        """
+        batch_img = []
+        ratios = []
+
+        for img in imgs:
+            preprocessed_img, ratio = super().preprocess(img)
+            batch_img.append(preprocessed_img)
+            ratios.append(ratio)
+
+        # Stack along the first dimension to create a batch
+        batch_img = np.stack(batch_img, axis=0)
+
+        return batch_img, ratios
+
+    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.
+        """
+        batch_img = batch_img.transpose(0, 3, 1, 2)  # NCHW format
+        batch_img = np.ascontiguousarray(batch_img, dtype=np.float32)
+
+        input = batch_img
+
+        # Create an IO Binding object
+        io_binding = self.session.io_binding()
+
+        # Bind the model inputs and outputs to the IO Binding object
+        io_binding.bind_input(name='input', device_type='cpu', device_id=0, element_type=np.float32, shape=input.shape, buffer_ptr=input.ctypes.data)
+        io_binding.bind_output(name='dets')
+        io_binding.bind_output(name='keypoints')
+
+        # Run inference with IO Binding
+        self.session.run_with_iobinding(io_binding)
+
+        # Retrieve the outputs from the IO Binding object
+        outputs = [output.numpy() for output in io_binding.get_outputs()]
+    
+        return outputs
+
+    def postprocess_batch(
+        self,
+        outputs: List[np.ndarray],
+        ratios: List[float]
+    ) -> List[Tuple[np.ndarray, 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.
+        """
+        batch_keypoints = []
+        batch_scores = []
+
+        for i, ratio in enumerate(ratios):
+            keypoints, scores = super().postprocess(outputs, ratio)
+            batch_keypoints.append(keypoints)
+            batch_scores.append(scores)
+
+        return batch_keypoints, batch_scores
+
+    def __call__(self, images: List[np.ndarray]):
+        batch_img, ratios = self.preprocess_batch(images)
+        outputs = self.inference(batch_img)
+        keypoints, scores = self.postprocess_batch(outputs, ratios)
+        return keypoints, scores