File size: 3,463 Bytes
dd253d4 | 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 | """Convert Stack-Can can+plate labeled tracks -> 3-PoinTr 2D dataset tree (label-conditioned).
Mirrors build_boxmark_3pointr.py.
Input: /data0/hongyi/reward_model_data/3pointr_data_stack_can_2d_canplate/<src>/tracks_labeled.npz
{tracks (T,N,2) px@384, vis (T,N), labels (N,) 0=can/1=plate, n_can, n_plate}
Output: /data0/hongyi/reward_model_data/3pointr_data_stack_can_canplate/no_actions/stack_can_canplate/<NNNN>/
timestep_k/{k_pos.pth (N,2)=tracks[0], k_cum_flow.pth (N,2)=tracks[k]-tracks[0],
k_visibility.pth (N,)}
labels.pth (N,) int64 (can first, then plate); metadata.json
EXCLUDES ep042 (false-positive plate: 23 pts, gap doesn't shrink). Tracks already 384x384.
"""
import os, json, glob
from pathlib import Path
import numpy as np
import torch
SRC = "/data0/hongyi/reward_model_data/3pointr_data_stack_can_2d_canplate"
OUT = "/data0/hongyi/reward_model_data/3pointr_data_stack_can_canplate/no_actions/stack_can_canplate"
EXCLUDE = {"tableHF_ep042"} # wrong-object plate
def main():
os.makedirs(OUT, exist_ok=True)
eps = sorted(glob.glob(os.path.join(SRC, "tableHF_ep*", "tracks_labeled.npz")))
print(f"[conv] {len(eps)} labeled-track episodes (excluding {sorted(EXCLUDE)})", flush=True)
written, minN, skipped = 0, 10**9, []
summary = []
for src_npz in eps:
src = Path(src_npz).parent.name
if src in EXCLUDE:
skipped.append(src); continue
d = np.load(src_npz)
tracks = d["tracks"].astype(np.float32) # (T,N,2)
labels = d["labels"].astype(np.int64) # (N,) 0=can 1=plate
n_can, n_plate = int(d["n_can"]), int(d["n_plate"])
if n_can == 0 or n_plate == 0: # safety: only both-object demos
skipped.append(src); continue
T, N, _ = tracks.shape
x0 = tracks[0] # (N,2)
out_dir = Path(OUT) / f"{written:04d}"
for k in range(T):
ts = out_dir / f"timestep_{k}"
ts.mkdir(parents=True, exist_ok=True)
torch.save(torch.from_numpy(x0.copy()), ts / f"{k}_pos.pth")
torch.save(torch.from_numpy((tracks[k] - x0).copy()), ts / f"{k}_cum_flow.pth")
torch.save(torch.from_numpy(np.ones(N, dtype=np.bool_)), ts / f"{k}_visibility.pth")
torch.save(torch.from_numpy(labels.copy()), out_dir / "labels.pth")
meta = {"T": int(T), "N": int(N), "point_dim": 2, "has_labels": True,
"n_can": n_can, "n_plate": n_plate, "src": src, "frame_size": [384, 384],
"intrinsics": {"camera": {"fx": 1.0, "fy": 1.0, "cx": 0.0, "cy": 0.0,
"width": 384, "height": 384}}}
with open(out_dir / "metadata.json", "w") as f:
json.dump(meta, f, indent=2)
summary.append({"idx": written, "src": src, "T": T, "N": N, "n_can": n_can, "n_plate": n_plate})
minN = min(minN, N)
written += 1
with open(os.path.join(os.path.dirname(OUT), "_canplate_build_summary.json"), "w") as f:
json.dump({"episodes": written, "min_N": minN, "skipped": skipped, "items": summary}, f, indent=2)
print(f"[conv] wrote {written} episodes -> {OUT} min total points N = {minN}", flush=True)
print(f"[conv] skipped: {skipped}", flush=True)
print(f"[conv] => set training --subsample <= {minN}", flush=True)
if __name__ == "__main__":
main()
|