| 100 | |
| 101 | |
| 102 | class TensorRTRefiner: |
| 103 | def __init__(self, args: argparse.Namespace): |
| 104 | import tensorrt as trt |
| 105 | |
| 106 | if not torch.cuda.is_available(): |
| 107 | raise RuntimeError("TensorRT backend requires CUDA") |
| 108 | model_path = Path(args.model) |
| 109 | if model_path.suffix != ".engine": |
| 110 | raise ValueError(f"--model must point to a TensorRT .engine file, got: {model_path}") |
| 111 | self.args = args |
| 112 | self.device = torch.device("cuda") |
| 113 | logger = trt.Logger(trt.Logger.WARNING) |
| 114 | runtime = trt.Runtime(logger) |
| 115 | self.engine = runtime.deserialize_cuda_engine(model_path.read_bytes()) |
| 116 | if self.engine is None: |
| 117 | raise RuntimeError(f"Failed to deserialize TensorRT engine: {model_path}") |
| 118 | self.context = self.engine.create_execution_context() |
| 119 | |
| 120 | names = [self.engine.get_tensor_name(i) for i in range(self.engine.num_io_tensors)] |
| 121 | input_names = [n for n in names if self.engine.get_tensor_mode(n) == trt.TensorIOMode.INPUT] |
| 122 | output_names = [n for n in names if self.engine.get_tensor_mode(n) == trt.TensorIOMode.OUTPUT] |
| 123 | self.image_name = "image" if "image" in input_names else input_names[0] |
| 124 | self.depth_name = "depth" if "depth" in input_names else input_names[1] |
| 125 | self.output_name = "depth_refined" if "depth_refined" in output_names else output_names[0] |
| 126 | self.image_shape = tuple(self.engine.get_tensor_shape(self.image_name)) |
| 127 | self.depth_shape = tuple(self.engine.get_tensor_shape(self.depth_name)) |
| 128 | self.output_shape = tuple(self.engine.get_tensor_shape(self.output_name)) |
| 129 | self.image_dtype = trt_dtype_to_torch(self.engine.get_tensor_dtype(self.image_name)) |
| 130 | self.depth_dtype = trt_dtype_to_torch(self.engine.get_tensor_dtype(self.depth_name)) |
| 131 | self.output_dtype = trt_dtype_to_torch(self.engine.get_tensor_dtype(self.output_name)) |
| 132 | self.stream = torch.cuda.Stream() |
| 133 | print(f"Loaded TensorRT engine: {args.model}") |
| 134 | print(f" {self.image_name}: {self.image_shape} {self.image_dtype}") |
| 135 | print(f" {self.depth_name}: {self.depth_shape} {self.depth_dtype}") |
| 136 | print(f" {self.output_name}: {self.output_shape} {self.output_dtype}") |
| 137 | |
| 138 | def refine(self, color_bgr: np.ndarray, depth_m: np.ndarray) -> tuple[np.ndarray, float]: |
| 139 | h, w = depth_m.shape |
| 140 | expected_h, expected_w = self.depth_shape[-2:] |
| 141 | if (h, w) != (expected_h, expected_w): |
| 142 | raise RuntimeError(f"Engine expects {expected_w}x{expected_h}, got {w}x{h}") |
| 143 | |
| 144 | image_rgb = cv2.cvtColor(color_bgr, cv2.COLOR_BGR2RGB) |
| 145 | image = torch.tensor(image_rgb / 255.0, dtype=self.image_dtype, device=self.device).permute(2, 0, 1).unsqueeze(0).contiguous() |
| 146 | depth = torch.tensor(depth_m, dtype=self.depth_dtype, device=self.device).unsqueeze(0).contiguous() |
| 147 | output = torch.empty(self.output_shape, dtype=self.output_dtype, device=self.device).contiguous() |
| 148 | |
| 149 | self.context.set_tensor_address(self.image_name, int(image.data_ptr())) |
| 150 | self.context.set_tensor_address(self.depth_name, int(depth.data_ptr())) |
| 151 | self.context.set_tensor_address(self.output_name, int(output.data_ptr())) |
| 152 | |
| 153 | start = time.perf_counter() |
| 154 | with torch.cuda.stream(self.stream): |
| 155 | ok = self.context.execute_async_v3(self.stream.cuda_stream) |
| 156 | self.stream.synchronize() |
| 157 | if not ok: |
| 158 | raise RuntimeError("TensorRT execution failed") |
| 159 | infer_s = time.perf_counter() - start |