| """Render a single grid PNG showing the first-frame RGB + sampled query points |
| (projected from 3D camera-frame back to pixels) for every episode in |
| 3pointr_data/no_actions/stack_can_50s_25f/0001..0075/. |
| |
| Missing episodes show a black panel labeled with the source video name. |
| """ |
| import argparse, json, os |
| from pathlib import Path |
|
|
| import numpy as np |
| import cv2 |
| import imageio |
|
|
|
|
| OUT_ROOT = "/home/hongyi/PWM/external/3pointr_data/no_actions/stack_can_50s_25f" |
|
|
|
|
| def video_label(E: int) -> str: |
| if 1 <= E <= 50: |
| return f"tableHF_ep{E-1:03d}" |
| if 51 <= E <= 75: |
| return f"table2_failure_ep{E+34:03d}" |
| return f"E={E}" |
|
|
|
|
| def project_pts(pts_xyz: np.ndarray, K: np.ndarray) -> np.ndarray: |
| """(N,3) cam-frame -> (N,2) pixel coords. Returns float.""" |
| fx, fy, cx, cy = K[0,0], K[1,1], K[0,2], K[1,2] |
| z = np.maximum(pts_xyz[:, 2], 1e-12) |
| u = pts_xyz[:, 0] / z * fx + cx |
| v = pts_xyz[:, 1] / z * fy + cy |
| return np.stack([u, v], axis=1) |
|
|
|
|
| def render_episode_panel(E: int, panel_h: int, panel_w: int) -> np.ndarray: |
| """Returns (panel_h, panel_w, 3) BGR image with overlay.""" |
| ep_dir = os.path.join(OUT_ROOT, f"{E:04d}") |
| rgb_path = os.path.join(ep_dir, "timestep_1", "1_rgb.png") |
| pos_path = os.path.join(ep_dir, "timestep_1", "1_pos.pth") |
| vis_path = os.path.join(ep_dir, "timestep_1", "1_visibility.pth") |
| meta_path = os.path.join(ep_dir, "metadata.json") |
|
|
| label = f"{E:02d} {video_label(E)}" |
|
|
| if not (os.path.exists(rgb_path) and os.path.exists(pos_path) and os.path.exists(meta_path)): |
| |
| panel = np.zeros((panel_h, panel_w, 3), dtype=np.uint8) |
| cv2.putText(panel, "MISSING", (8, 30), |
| cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2, cv2.LINE_AA) |
| cv2.putText(panel, label, (8, panel_h - 10), |
| cv2.FONT_HERSHEY_SIMPLEX, 0.45, (255, 255, 255), 1, cv2.LINE_AA) |
| return panel |
|
|
| rgb = imageio.imread(rgb_path) |
| H, W = rgb.shape[:2] |
| img = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR) |
|
|
| |
| import torch |
| pts = torch.load(pos_path, map_location="cpu", weights_only=False).numpy() |
| vis = None |
| if os.path.exists(vis_path): |
| vis = torch.load(vis_path, map_location="cpu", weights_only=False).numpy().astype(bool) |
| meta = json.load(open(meta_path)) |
| cam = next(iter(meta["intrinsics"].values())) |
| K = np.array([[cam["fx"], 0, cam["cx"]], |
| [0, cam["fy"], cam["cy"]], |
| [0, 0, 1]], dtype=np.float32) |
|
|
| uv = project_pts(pts, K) |
| valid = (uv[:,0] >= 0) & (uv[:,0] < W) & (uv[:,1] >= 0) & (uv[:,1] < H) & (pts[:,2] > 1e-6) |
| if vis is not None: |
| valid_vis = valid & vis |
| valid_invis = valid & (~vis) |
| else: |
| valid_vis = valid |
| valid_invis = np.zeros_like(valid) |
|
|
| |
| for u, v in uv[valid_invis].astype(int): |
| cv2.circle(img, (u, v), 1, (128, 128, 128), -1, cv2.LINE_AA) |
| for u, v in uv[valid_vis].astype(int): |
| cv2.circle(img, (u, v), 1, (0, 0, 255), -1, cv2.LINE_AA) |
|
|
| |
| panel = cv2.resize(img, (panel_w, panel_h), interpolation=cv2.INTER_AREA) |
|
|
| n_vis = int(valid_vis.sum()); n_total = pts.shape[0] |
| cv2.putText(panel, label, (4, panel_h - 22), |
| cv2.FONT_HERSHEY_SIMPLEX, 0.40, (0, 0, 0), 3, cv2.LINE_AA) |
| cv2.putText(panel, label, (4, panel_h - 22), |
| cv2.FONT_HERSHEY_SIMPLEX, 0.40, (255, 255, 255), 1, cv2.LINE_AA) |
| cv2.putText(panel, f"vis {n_vis}/{n_total}", (4, panel_h - 6), |
| cv2.FONT_HERSHEY_SIMPLEX, 0.40, (0, 0, 0), 3, cv2.LINE_AA) |
| cv2.putText(panel, f"vis {n_vis}/{n_total}", (4, panel_h - 6), |
| cv2.FONT_HERSHEY_SIMPLEX, 0.40, (255, 255, 255), 1, cv2.LINE_AA) |
| return panel |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--cols", type=int, default=10) |
| ap.add_argument("--panel_h", type=int, default=192) |
| ap.add_argument("--panel_w", type=int, default=192) |
| ap.add_argument("--out", default="/home/hongyi/PWM/external/3pointr_data/stack_can_50s_25f_first_frames.png") |
| args = ap.parse_args() |
|
|
| cols = int(args.cols) |
| rows = int(np.ceil(75 / cols)) |
| grid = np.zeros((rows * args.panel_h, cols * args.panel_w, 3), dtype=np.uint8) |
|
|
| for E in range(1, 76): |
| r = (E - 1) // cols |
| c = (E - 1) % cols |
| panel = render_episode_panel(E, args.panel_h, args.panel_w) |
| grid[r*args.panel_h:(r+1)*args.panel_h, c*args.panel_w:(c+1)*args.panel_w] = panel |
|
|
| cv2.imwrite(args.out, grid) |
| print(f"wrote {args.out} grid={rows}x{cols} size={grid.shape[1]}x{grid.shape[0]}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
