| 99 | |
| 100 | |
| 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.') |
| 121 | |
| 122 | def encode_decode(self, img, img_metas): |
| 123 | raise NotImplementedError('This method is not implemented.') |
| 124 | |
| 125 | def forward_train(self, imgs, img_metas, **kwargs): |
| 126 | raise NotImplementedError('This method is not implemented.') |
| 127 | |
| 128 | def simple_test(self, img: torch.Tensor, img_meta: Iterable, |
| 129 | **kwargs) -> list: |
| 130 | with torch.cuda.device(self.device_id), torch.no_grad(): |
| 131 | seg_pred = self.model({'input': img})['output'] |
| 132 | seg_pred = seg_pred.detach().cpu().numpy() |
| 133 | # whole might support dynamic reshape |
| 134 | ori_shape = img_meta[0]['ori_shape'] |
| 135 | if not (ori_shape[0] == seg_pred.shape[-2] |
| 136 | and ori_shape[1] == seg_pred.shape[-1]): |
| 137 | seg_pred = torch.from_numpy(seg_pred).float() |
| 138 | seg_pred = resize( |
| 139 | seg_pred, size=tuple(ori_shape[:2]), mode='nearest') |
| 140 | seg_pred = seg_pred.long().detach().cpu().numpy() |
| 141 | seg_pred = seg_pred[0] |
| 142 | seg_pred = list(seg_pred) |
| 143 | return seg_pred |
| 144 | |
| 145 | def aug_test(self, imgs, img_metas, **kwargs): |
| 146 | raise NotImplementedError('This method is not implemented.') |
| 147 | |
| 148 | |
| 149 | def parse_args() -> argparse.Namespace: |