Parse raw NPZ arrays into the standardized format. Handles both single-file and per-frame-dir data (same key structure). Returns dict: images (S,H,W,3) uint8, depth (S,H,W) float32, c2w (S,4,4) float32, K (S,3,3) float32, confidence (S,H,W) float32 or No
(data: Dict[str, np.ndarray])
| 68 | |
| 69 | |
| 70 | def _parse_raw_data(data: Dict[str, np.ndarray]) -> Dict: |
| 71 | """Parse raw NPZ arrays into the standardized format. |
| 72 | |
| 73 | Handles both single-file and per-frame-dir data (same key structure). |
| 74 | |
| 75 | Returns dict: images (S,H,W,3) uint8, depth (S,H,W) float32, |
| 76 | c2w (S,4,4) float32, K (S,3,3) float32, |
| 77 | confidence (S,H,W) float32 or None. |
| 78 | """ |
| 79 | images = data['images'] |
| 80 | if images.ndim == 4 and images.shape[1] == 3: |
| 81 | images = np.ascontiguousarray(images.transpose(0, 2, 3, 1)) |
| 82 | if images.dtype != np.uint8: |
| 83 | images = (images * 255).clip(0, 255).astype(np.uint8) if images.max() <= 1.0 else images.astype(np.uint8) |
| 84 | |
| 85 | depth = data['depth'].astype(np.float32) |
| 86 | if depth.ndim == 4: |
| 87 | depth = depth[..., 0] |
| 88 | |
| 89 | K_raw = data['intrinsic'].astype(np.float32) |
| 90 | if K_raw.ndim == 2: |
| 91 | K_raw = np.tile(K_raw[None], (len(images), 1, 1)) |
| 92 | |
| 93 | if 'extrinsic' not in data: |
| 94 | raise ValueError("NPZ must contain 'extrinsic' (W2C poses).") |
| 95 | ext = data['extrinsic'].astype(np.float32) |
| 96 | nf = ext.shape[0] |
| 97 | w2c = np.zeros((nf, 4, 4), dtype=np.float32) |
| 98 | w2c[:, :3, :] = ext[:, :3, :] |
| 99 | w2c[:, 3, 3] = 1.0 |
| 100 | R = w2c[:, :3, :3] |
| 101 | t = w2c[:, :3, 3:4] |
| 102 | Rt = R.transpose(0, 2, 1) |
| 103 | c2w = np.zeros((nf, 4, 4), dtype=np.float32) |
| 104 | c2w[:, :3, :3] = Rt |
| 105 | c2w[:, :3, 3:4] = -Rt @ t |
| 106 | c2w[:, 3, 3] = 1.0 |
| 107 | |
| 108 | confidence = None |
| 109 | for conf_key in ('depth_conf', 'confidence'): |
| 110 | if conf_key in data: |
| 111 | confidence = data[conf_key].astype(np.float32) |
| 112 | if confidence.ndim == 4: |
| 113 | confidence = confidence[..., 0] |
| 114 | break |
| 115 | |
| 116 | # Optional keyframe mask from meta.npz (produced by batch_demo.py). |
| 117 | # ``is_keyframe`` is the preferred key (bool per frame); ``frame_type`` |
| 118 | # (uint8, 0=scale, 1=keyframe, 2=non-keyframe) is accepted as a fallback. |
| 119 | is_keyframe = None |
| 120 | if 'is_keyframe' in data: |
| 121 | is_keyframe = np.asarray(data['is_keyframe']).astype(bool) |
| 122 | elif 'frame_type' in data: |
| 123 | is_keyframe = np.asarray(data['frame_type']) != 2 |
| 124 | if is_keyframe is not None: |
| 125 | is_keyframe = np.squeeze(is_keyframe) |
| 126 | if is_keyframe.ndim == 0: |
| 127 | is_keyframe = is_keyframe.reshape(1) |