File size: 4,828 Bytes
23e1f91 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | """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)):
# MISSING — black panel with label
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, 3) RGB
H, W = rgb.shape[:2]
img = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
# Load 3D pos and intrinsics, project, overlay
import torch
pts = torch.load(pos_path, map_location="cpu", weights_only=False).numpy() # (N,3)
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) # (N,2)
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)
# Draw invisible points (gray) first, then visible (red) on top
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)
# Resize to panel size
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()
|