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