| """Convert box+mark labeled tracks -> 3-PoinTr 2D dataset tree, with per-episode labels. |
| |
| Input: /data0/hongyi/reward_model_data/3pointr_data_push_box_2d_boxmark/<src>/tracks_labeled.npz |
| {tracks (T,N,2) px@384, vis (T,N), labels (N,) 0=box/1=mark, n_box, n_mark} |
| Output: /data0/hongyi/reward_model_data/3pointr_data_push_box_boxmark/no_actions/push_box_boxmark/<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 <- box/mark, aligned with point order (box first, then mark) |
| metadata.json {point_dim:2, has_labels:true, n_box, n_mark} |
| |
| Points/flow stay 2D (normalization untouched); the label is consumed by a label-embedding |
| in the model, NOT as a coordinate. Tracks are already 384x384 (videos_small) -> no rescale. |
| ignore_vis=True (trust CoTracker through-occlusion positions; box is occluded by hand/mark). |
| """ |
| import os, json, glob |
| from pathlib import Path |
| import numpy as np |
| import torch |
|
|
| SRC = "/data0/hongyi/reward_model_data/3pointr_data_push_box_2d_boxmark" |
| OUT = "/data0/hongyi/reward_model_data/3pointr_data_push_box_boxmark/no_actions/push_box_boxmark" |
|
|
|
|
| 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", flush=True) |
| written, minN = 0, 10**9 |
| summary = [] |
| for src_npz in eps: |
| src = Path(src_npz).parent.name |
| d = np.load(src_npz) |
| tracks = d["tracks"].astype(np.float32) |
| labels = d["labels"].astype(np.int64) |
| n_box, n_mark = int(d["n_box"]), int(d["n_mark"]) |
| 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_box": n_box, "n_mark": n_mark, "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_box": n_box, "n_mark": n_mark}) |
| minN = min(minN, N) |
| written += 1 |
| with open(os.path.join(os.path.dirname(OUT), "_boxmark_build_summary.json"), "w") as f: |
| json.dump({"episodes": written, "min_N": minN, "items": summary}, f, indent=2) |
| print(f"[conv] wrote {written} episodes -> {OUT} min total points N = {minN}", flush=True) |
| print(f"[conv] => set training --subsample <= {minN}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
