0

I'm new in onnx. How I can convert yolov5 to onnx and get bbox?

class YOLOv5ONNX:
    def __init__(self, model_file: str, device: str = "cpu"):
        providers = [("CPUExecutionProvider")]

        self.onnx_session = onnxruntime.InferenceSession(
            model_file, providers=providers
        )

    def __preprocess(self, image: np.ndarray):
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        image = cv2.resize(image, (640, 640))
        image = image.astype(np.float32) / 255.0
        image = np.transpose(image, axes=(2, 0, 1))
        image = np.expand_dims(image, 0)
        
        return image

    def __call__(self, image: np.ndarray):
        work_image = self.__preprocess(image)
        result = self.onnx_session.run(
            None, {self.onnx_session.get_inputs()[0].name: work_image}
        )
        return result

Output: https://i.stack.imgur.com/WwgP4.png

Bombex
  • 1
  • 2

0 Answers0