| 199 | |
| 200 | |
| 201 | def load_capture(capture_dir: Path, device: torch.device) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: |
| 202 | rgb_path = capture_dir / "rgb.png" |
| 203 | depth_path = capture_dir / "raw_depth.png" |
| 204 | metadata_path = capture_dir / "metadata.json" |
| 205 | bgr = cv2.imread(str(rgb_path), cv2.IMREAD_COLOR) |
| 206 | depth_mm = cv2.imread(str(depth_path), cv2.IMREAD_UNCHANGED) |
| 207 | if bgr is None: |
| 208 | raise FileNotFoundError(f"Failed to read RGB image: {rgb_path}") |
| 209 | if depth_mm is None: |
| 210 | raise FileNotFoundError(f"Failed to read depth image: {depth_path}") |
| 211 | |
| 212 | rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) |
| 213 | image = torch.tensor(rgb / 255.0, dtype=torch.float32, device=device).permute(2, 0, 1).unsqueeze(0).contiguous() |
| 214 | depth = torch.tensor(depth_mm.astype(np.float32) / 1000.0, dtype=torch.float32, device=device).unsqueeze(0).contiguous() |
| 215 | |
| 216 | metadata: dict[str, Any] = {} |
| 217 | if metadata_path.exists(): |
| 218 | metadata = json.loads(metadata_path.read_text(encoding="utf-8")) |
| 219 | metadata.update({"height": int(rgb.shape[0]), "width": int(rgb.shape[1])}) |
| 220 | return image, depth, metadata |
| 221 | |
| 222 | |
| 223 | def make_dummy_inputs(width: int, height: int, device: torch.device) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: |