class ONNXModel: def init(self, model_path): self.model_path = model_path self.model = None

def load_model(self):
    pass

def predict(self, inputs):
    pass

class ResNet50Model(ONNXModel): def load_model(self): import onnxruntime self.model = onnxruntime.InferenceSession(self.model_path)

def predict(self, inputs):
    import numpy as np
    inputs = np.array(inputs).astype(np.float32)
    outputs = self.model.run(None, {'input': inputs})
    return outputs[0]

class YOLOv3Model(ONNXModel): def load_model(self): import onnxruntime self.model = onnxruntime.InferenceSession(self.model_path)

def predict(self, inputs):
    import numpy as np
    inputs = np.array(inputs).astype(np.float32)
    outputs = self.model.run(None, {'input': inputs})
    return outputs[0]

使用方法

if name == 'main': resnet_model = ResNet50Model('resnet50.onnx') resnet_model.load_model() resnet_output = resnet_model.predict([1, 2, 3])

yolov3_model = YOLOv3Model('yolov3.onnx')
yolov3_model.load_model()
yolov3_output = yolov3_model.predict([1, 2, 3])

print(resnet_output)
print(yolov3_output)