| """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"} |
|
|
|
|
| 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) |
| labels = d["labels"].astype(np.int64) |
| n_can, n_plate = int(d["n_can"]), int(d["n_plate"]) |
| if n_can == 0 or n_plate == 0: |
| skipped.append(src); continue |
| T, N, _ = tracks.shape |
| x0 = tracks[0] |
| 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() |
|
|
