| 101 | class TensorRTSegmentor(BaseSegmentor): |
| 102 | |
| 103 | def __init__(self, trt_file: str, cfg: Any, device_id: int): |
| 104 | super(TensorRTSegmentor, self).__init__() |
| 105 | from mmcv.tensorrt import TRTWraper, load_tensorrt_plugin |
| 106 | try: |
| 107 | load_tensorrt_plugin() |
| 108 | except (ImportError, ModuleNotFoundError): |
| 109 | warnings.warn('If input model has custom op from mmcv, \ |
| 110 | you may have to build mmcv with TensorRT from source.') |
| 111 | model = TRTWraper( |
| 112 | trt_file, input_names=['input'], output_names=['output']) |
| 113 | |
| 114 | self.model = model |
| 115 | self.device_id = device_id |
| 116 | self.cfg = cfg |
| 117 | self.test_mode = cfg.model.test_cfg.mode |
| 118 | |
| 119 | def extract_feat(self, imgs): |
| 120 | raise NotImplementedError('This method is not implemented.') |