Create physical_dataset.py

physical_dataset.py

@@ -0,0 +1,112 @@
+"""
+physical_dataset.py
+
+Dataset for Contact / Force VAE finetuning.
+
+Reads clips.json (from build_clip_index.py) and statistics.json, loads each
+17-frame clip of per-finger contact (2,H,W) or force (6,H,W) maps, normalizes,
+and returns a tensor plus an active mask for weighted reconstruction loss.
+
+Normalization (background stays 0 in both cases):
+  contact: x[ch] / contact_ch_max[ch]          -> (0, 1], background 0
+  force:   x[ch] / force_ch_std[ch]            -> ~O(1) signed, background 0
+Rationale: contact/force are sparse; subtracting a mean would turn the 0
+background into a nonzero value and destroy sparsity, so we only scale.
+
+Returns per item:
+  {
+    "data":        (C, T, H, W) float32, normalized
+    "active_mask": (1, T, H, W) float32, 1 where any channel != 0
+    "episode":     str
+  }
+"""
+
+import json
+import os
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+
+class PhysicalClipDataset(Dataset):
+    def __init__(self, clips_json, statistics_json, source_root,
+                 modality="contact", eps=1e-6):
+        assert modality in ("contact", "force")
+        self.modality = modality
+        self.source_root = source_root
+        self.eps = eps
+
+        with open(clips_json) as f:
+            blob = json.load(f)
+        self.clips = blob["clips"] if isinstance(blob, dict) and "clips" in blob else blob
+        self.config = blob.get("config", {}) if isinstance(blob, dict) else {}
+
+        with open(statistics_json) as f:
+            self.stats = json.load(f)
+
+        # build per-channel scale vector, shaped (C,1,1,1) for broadcasting
+        if modality == "contact":
+            scale = np.asarray(self.stats["contact_ch_max"], dtype=np.float32)  # (2,)
+        else:
+            scale = np.asarray(self.stats["force_ch_std"], dtype=np.float32)     # (6,)
+        scale = np.maximum(scale, eps)                # avoid div-by-zero
+        self.scale = scale.reshape(-1, 1, 1, 1)       # (C,1,1,1)
+        self.n_ch = self.scale.shape[0]
+
+        self.path_key = "contact_paths" if modality == "contact" else "force_paths"
+
+    def __len__(self):
+        return len(self.clips)
+
+    def _load_clip(self, clip):
+        """Load 17 frames -> (C, T, H, W) raw float32."""
+        ep = clip["episode"]
+        frames = []
+        for rel in clip[self.path_key]:
+            arr = np.load(os.path.join(self.source_root, ep, rel))  # (C,H,W)
+            frames.append(arr.astype(np.float32))
+        # stack along time: list of (C,H,W) -> (T,C,H,W) -> (C,T,H,W)
+        x = np.stack(frames, axis=0).transpose(1, 0, 2, 3)
+        return x
+
+    def __getitem__(self, idx):
+        clip = self.clips[idx]
+        x = self._load_clip(clip)                    # (C,T,H,W) raw
+
+        # active mask BEFORE normalization (any channel nonzero)
+        active = (np.abs(x).sum(axis=0, keepdims=True) > 0).astype(np.float32)  # (1,T,H,W)
+
+        # normalize: scale only, background 0 stays 0
+        x = x / self.scale                           # (C,T,H,W)
+
+        return {
+            "data": torch.from_numpy(x),             # (C,T,H,W)
+            "active_mask": torch.from_numpy(active),  # (1,T,H,W)
+            "episode": clip["episode"],
+        }
+
+
+def collate_physical(batch):
+    """Stack into (B,C,T,H,W). Assumes uniform shape (it is: fixed 17 frames)."""
+    data = torch.stack([b["data"] for b in batch], dim=0)
+    mask = torch.stack([b["active_mask"] for b in batch], dim=0)
+    episodes = [b["episode"] for b in batch]
+    return {"data": data, "active_mask": mask, "episodes": episodes}
+
+
+if __name__ == "__main__":
+    import argparse
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--clips", required=True)
+    ap.add_argument("--stats", required=True)
+    ap.add_argument("--source_root", required=True)
+    ap.add_argument("--modality", choices=["contact", "force"], default="contact")
+    args = ap.parse_args()
+
+    ds = PhysicalClipDataset(args.clips, args.stats, args.source_root, args.modality)
+    print(f"{args.modality} dataset: {len(ds)} clips, scale={ds.scale.ravel()}")
+    item = ds[0]
+    print("data:", tuple(item["data"].shape), item["data"].dtype,
+          "min/max:", float(item["data"].min()), float(item["data"].max()))
+    print("active_mask:", tuple(item["active_mask"].shape),
+          "active frac:", float(item["active_mask"].mean()))