(self, args: argparse.Namespace)
| 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 |
nothing calls this directly
no test coverage detected