| 16 | # from cv2box.utils import try_import |
| 17 | |
| 18 | class ModelBase: |
| 19 | def __init__(self, model_info, provider): |
| 20 | self.model_path = model_info['model_path'] |
| 21 | |
| 22 | if 'input_dynamic_shape' in model_info.keys(): |
| 23 | self.input_dynamic_shape = model_info['input_dynamic_shape'] |
| 24 | else: |
| 25 | self.input_dynamic_shape = None |
| 26 | |
| 27 | if 'picklable' in model_info.keys(): |
| 28 | picklable = model_info['picklable'] |
| 29 | else: |
| 30 | picklable = False |
| 31 | |
| 32 | if 'trt_wrapper_self' in model_info.keys(): |
| 33 | TRTWrapper = TRTWrapperSelf |
| 34 | |
| 35 | # init model |
| 36 | if Path(self.model_path).suffix == '.engine': |
| 37 | self.model_type = 'trt' |
| 38 | self.model = TRTWrapper(self.model_path) |
| 39 | elif Path(self.model_path).suffix == '.tjm': |
| 40 | self.model_type = 'tjm' |
| 41 | self.model = TJMWrapper(self.model_path, provider=provider) |
| 42 | elif Path(self.model_path).suffix in ['.onnx', '.bin']: |
| 43 | self.model_type = 'onnx' |
| 44 | if not picklable: |
| 45 | if 'encrypt' in model_info.keys(): |
| 46 | self.model_path = load_encrypt_model(self.model_path, key=model_info['encrypt']) |
| 47 | self.model = ONNXModel(self.model_path, provider=provider, input_dynamic_shape=self.input_dynamic_shape) |
| 48 | else: |
| 49 | self.model = OnnxModelPickable(self.model_path, provider=provider, ) |
| 50 | else: |
| 51 | raise 'check model suffix , support engine/tjm/onnx now.' |
nothing calls this directly
no outgoing calls
no test coverage detected