| 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 |
| 160 | return output.squeeze(0).detach().float().cpu().numpy(), infer_s |
| 161 | |
| 162 | |
| 163 | def open_video_writer(args: argparse.Namespace, frame_size: tuple[int, int]) -> cv2.VideoWriter | None: |