Buckets:
| """Pre-render mesh+hand overlays for all TACO sequences. | |
| Uses the same rendering approach as notebooks/explore_taco.ipynb (cell 18). | |
| Usage: | |
| .venv/bin/python scripts/prerender_overlays.py | |
| Outputs: outputs/overlays/<sequence_id_sanitized>.jpg | |
| """ | |
| import sys | |
| import numpy as np | |
| import pandas as pd | |
| import cv2 | |
| import torch | |
| from pathlib import Path | |
| PROJECT_ROOT = Path(__file__).resolve().parent.parent | |
| sys.path.insert(0, str(PROJECT_ROOT)) | |
| import decord | |
| decord.bridge.set_bridge("torch") | |
| from fwomo_3d.config import TacoDatasetCfg, ViewSamplerAllCfg | |
| from fwomo_3d.data.taco_utils import load_object_meshes, render_mesh_on_image | |
| from fwomo_3d.utils.visualization.hand import ( | |
| HAND_COLORS, | |
| c2w_and_norm_K_to_w2c, | |
| draw_hand_skeleton, | |
| ) | |
| TACO = Path("/mnt/fast/mzhobro/taco_dataset_resized") | |
| OUT = PROJECT_ROOT / "outputs" / "overlays" | |
| OUT.mkdir(parents=True, exist_ok=True) | |
| GRID_ROWS, GRID_COLS = 3, 4 | |
| CAM_W, CAM_H = 384, 282 | |
| HEADER_H = 36 | |
| def sanitize(seq_id): | |
| return seq_id.replace("/", "__").replace(" ", "_").replace("(", "").replace(")", "").replace(",", "") | |
| def main(): | |
| # Build dataset (same config as explore_taco notebook) | |
| cfg = TacoDatasetCfg( | |
| root=str(TACO), | |
| image_shape=[376, 512], | |
| view_sampler=ViewSamplerAllCfg(name="all"), | |
| near=0.1, | |
| far=5.0, | |
| load_hand_poses=True, | |
| load_object_poses=True, | |
| ) | |
| # Build all splits to cover every sequence | |
| ds_train = cfg.build("train") | |
| ds_val = cfg.build("val") | |
| ds_test = cfg.build("test") | |
| # Merge metadata from all splits | |
| ds = ds_train | |
| all_meta = pd.concat([ds_train.meta, ds_val.meta, ds_test.meta], ignore_index=True) | |
| ds.meta = all_meta | |
| models_dir = TACO / "object_models_released" | |
| n = len(ds.meta) | |
| existing = len(list(OUT.glob("*.jpg"))) | |
| print(f"Rendering {n} sequences ({existing} already done)") | |
| for seq_idx in range(n): | |
| row = ds.meta.iloc[seq_idx] | |
| seq_id = row["sequence_id"] | |
| out_path = OUT / f"{sanitize(seq_id)}.jpg" | |
| if out_path.exists(): | |
| continue | |
| try: | |
| frame_idx = 0 | |
| views = ds.load_sequence_views(seq_idx, frame_idx) | |
| H, W = views["image"].shape[2], views["image"].shape[3] | |
| n_cams = len(views["index"]) | |
| # Load meshes | |
| meshes = load_object_meshes(TACO, row, frame_idx, models_dir) | |
| # Load hands from actions | |
| hands = {} | |
| if "actions" in views: | |
| for key in ("left_hand", "right_hand"): | |
| if key in views["actions"]: | |
| joints = views["actions"][key][0].numpy() # (21, 3) | |
| if not np.allclose(joints, 0.0): | |
| side = key.replace("_hand", "") | |
| hands[side] = joints.astype(np.float64) | |
| # Build grid | |
| canvas = np.zeros((GRID_ROWS * CAM_H + HEADER_H, GRID_COLS * CAM_W, 3), dtype=np.uint8) | |
| for v in range(min(n_cams, GRID_ROWS * GRID_COLS)): | |
| img_np = (views["image"][v].permute(1, 2, 0).numpy() * 255).astype(np.uint8).copy() | |
| c2w = views["extrinsics"][v].numpy() | |
| K_norm = views["intrinsics"][v].numpy() | |
| R, T, K_px = c2w_and_norm_K_to_w2c(c2w, K_norm, H, W) | |
| # Render meshes | |
| for verts_w, faces_arr, color, label in meshes: | |
| img_np = render_mesh_on_image(img_np, verts_w, faces_arr, K_px, R, T, color=color) | |
| # Render hands | |
| for side, joints in hands.items(): | |
| img_np = draw_hand_skeleton(img_np, joints, K_px, R, T, HAND_COLORS[side]) | |
| # Camera label | |
| cam_label = views["index"][v] if isinstance(views["index"][v], str) else str(views["index"][v]) | |
| cv2.putText(img_np, cam_label[-4:], (5, 22), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2) | |
| # Place in grid | |
| r, c = v // GRID_COLS, v % GRID_COLS | |
| y0, x0 = r * CAM_H + HEADER_H, c * CAM_W | |
| resized = cv2.resize(img_np, (CAM_W, CAM_H)) | |
| canvas[y0:y0 + CAM_H, x0:x0 + CAM_W] = resized | |
| # Header | |
| cv2.rectangle(canvas, (0, 0), (canvas.shape[1], HEADER_H - 2), (40, 40, 40), -1) | |
| header = f"{seq_id} (date={row.get('date', '?')})" | |
| cv2.putText(canvas, header, (10, 26), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1) | |
| # Save as BGR | |
| cv2.imwrite(str(out_path), cv2.cvtColor(canvas, cv2.COLOR_RGB2BGR), | |
| [cv2.IMWRITE_JPEG_QUALITY, 85]) | |
| except Exception as e: | |
| print(f" ERROR {seq_id}: {e}") | |
| if (seq_idx + 1) % 50 == 0: | |
| done = len(list(OUT.glob("*.jpg"))) | |
| print(f" {seq_idx + 1}/{n} done ({done} total)") | |
| print(f"Done. {len(list(OUT.glob('*.jpg')))} overlays in {OUT}") | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 5.01 kB
- Xet hash:
- fd04c1f06e12e882480808f54a0354f140d9ed8a34f33efd3923cf0af8efba03
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.