robocandywrapper 0.2.14__tar.gz → 0.2.16__tar.gz

{robocandywrapper-0.2.14/robocandywrapper.egg-info → robocandywrapper-0.2.16}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: robocandywrapper
-Version: 0.2.14
+Version: 0.2.16
 Summary: Sweet wrappers for extending and remixing LeRobot Datasets
 Author: RoboCandyWrapper Contributors
 License: MIT License
@@ -43,13 +43,14 @@ Requires-Dist: torch>=2.0.0
 Requires-Dist: lerobot<0.5,>=0.4
 Requires-Dist: pandas>=1.3.0
 Requires-Dist: rewact_tools
+Requires-Dist: motion-primitives
+Requires-Dist: lerobot-policy-diffusion-motion-primitives
+Requires-Dist: lerobot-policy-diffusion-pointing
 Provides-Extra: dev
 Requires-Dist: pytest>=7.0.0; extra == "dev"
 Requires-Dist: black>=22.0.0; extra == "dev"
 Requires-Dist: isort>=5.10.0; extra == "dev"
 Requires-Dist: flake8>=4.0.0; extra == "dev"
-Requires-Dist: motion-primitives; extra == "dev"
-Requires-Dist: lerobot-policy-diffusion-motion-primitives; extra == "dev"
 Dynamic: license-file
 Dynamic: requires-python
 

{robocandywrapper-0.2.14 → robocandywrapper-0.2.16}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "robocandywrapper"
-version = "0.2.14"
+version = "0.2.16"
 description = "Sweet wrappers for extending and remixing LeRobot Datasets"
 readme = "README.md"
 requires-python = ">=3.10,<3.11"
@@ -28,6 +28,9 @@ dependencies = [
     "lerobot>=0.4,<0.5",
     "pandas>=1.3.0",
     "rewact_tools",
+    "motion-primitives",
+    "lerobot-policy-diffusion-motion-primitives",
+    "lerobot-policy-diffusion-pointing",
 ]
 
 [project.optional-dependencies]
@@ -36,8 +39,6 @@ dev = [
     "black>=22.0.0",
     "isort>=5.10.0",
     "flake8>=4.0.0",
-    "motion-primitives",
-    "lerobot-policy-diffusion-motion-primitives"
 ]
 
 [project.urls]
@@ -57,6 +58,7 @@ override-dependencies = [
 
 [tool.uv.sources]
 lerobot-policy-diffusion-motion-primitives = { path = "../../experimental/lerobot_policy_diffusion_motion_primitives", editable = true }
+lerobot-policy-diffusion-pointing = { path = "../../experimental/lerobot_policy_diffusion_pointing", editable = true }
 motion-primitives = { path = "../../experimental/motion_primitives", editable = true }
 rewact-tools = { path = "../rewACT/rewact_tools", editable = true }
 

{robocandywrapper-0.2.14 → robocandywrapper-0.2.16}/robocandywrapper/__init__.py

@@ -18,7 +18,7 @@ from robocandywrapper.constants import (
     MOTION_PRIMITIVE_PLUGIN_NAME,
 )
 
-__version__ = "0.2.14"
+__version__ = "0.2.16"
 
 __all__ = [
     "DatasetPlugin",

{robocandywrapper-0.2.14 → robocandywrapper-0.2.16}/robocandywrapper/constants.py

@@ -10,5 +10,6 @@ AFFORDANCE_PLUGIN_NAME = "affordance"
 CONTROL_MODE_PLUGIN_NAME = "control_mode"
 EPISODE_OUTCOME_PLUGIN_NAME = "episode_outcome"
 SUBTASK_PLUGIN_NAME = "subtask"
+MOLMOPOINT_PLUGIN_NAME = "molmopoint"
 MOTION_PRIMITIVE_PLUGIN_NAME = "motion_primitives"
 

{robocandywrapper-0.2.14 → robocandywrapper-0.2.16}/robocandywrapper/dataformats/lerobot_21/convert_v20_to_v21.py

@@ -34,7 +34,7 @@ from huggingface_hub import HfApi
 from tqdm import tqdm
 
 from lerobot.datasets.compute_stats import aggregate_stats, get_feature_stats, sample_indices
-from lerobot.datasets.utils import load_stats, write_info
+from lerobot.datasets.io_utils import load_stats, write_info
 
 from robocandywrapper.dataformats.lerobot_21.dataset import LeRobot21Dataset
 from robocandywrapper.dataformats.lerobot_21.utils import (

{robocandywrapper-0.2.14 → robocandywrapper-0.2.16}/robocandywrapper/dataformats/lerobot_21/dataset.py

@@ -38,24 +38,34 @@ from lerobot.datasets.utils import (
     DEFAULT_IMAGE_PATH,
     INFO_PATH,
     LEGACY_TASKS_PATH as TASKS_PATH,
+    check_version_compatibility,
+    create_lerobot_dataset_card,
+    get_safe_version,
+    is_valid_version,
+)
+from lerobot.datasets.feature_utils import (
     _validate_feature_names,
     check_delta_timestamps,
-    check_version_compatibility,
     create_empty_dataset_info,
-    create_lerobot_dataset_card,
-    embed_images,
     get_delta_indices,
     get_hf_features_from_features,
-    get_safe_version,
-    hf_transform_to_torch,
-    is_valid_version,
-    load_info,
-    load_stats,
     validate_episode_buffer,
     validate_frame,
+)
+from lerobot.datasets.io_utils import (
+    embed_images,
+    load_info,
+    load_stats,
     write_info,
     write_json,
 )
+
+def hf_transform_to_torch(item):
+    """Compat shim: transforms HF dataset items to torch tensors."""
+    for key in item:
+        if isinstance(item[key], np.ndarray):
+            item[key] = torch.from_numpy(item[key])
+    return item
 # Import legacy v2.1-specific functions that were removed in lerobot 0.4.1
 from robocandywrapper.dataformats.lerobot_21.utils import (
     append_jsonlines,
@@ -749,7 +759,7 @@ class LeRobot21Dataset(torch.utils.data.Dataset):
 
     def __getitem__(self, idx) -> dict:
         item = self.hf_dataset[idx]
-        ep_idx = item["episode_index"].item()
+        ep_idx = item["episode_index"].item() if isinstance(item["episode_index"], torch.Tensor) else item["episode_index"]
 
         query_indices = None
         if self.delta_indices is not None:
@@ -760,7 +770,7 @@ class LeRobot21Dataset(torch.utils.data.Dataset):
                 item[key] = val
 
         if len(self.meta.video_keys) > 0:
-            current_ts = item["timestamp"].item()
+            current_ts = item["timestamp"].item() if isinstance(item["timestamp"], torch.Tensor) else item["timestamp"]
             query_timestamps = self._get_query_timestamps(current_ts, query_indices)
             video_frames = self._query_videos(query_timestamps, ep_idx)
             item = {**video_frames, **item}
@@ -771,7 +781,7 @@ class LeRobot21Dataset(torch.utils.data.Dataset):
                 item[cam] = self.image_transforms(item[cam])
 
         # Add task as a string
-        task_idx = item["task_index"].item()
+        task_idx = item["task_index"].item() if isinstance(item["task_index"], torch.Tensor) else item["task_index"]
         item["task"] = self.meta.tasks[task_idx]
 
         # Hack - add gripper position to end

{robocandywrapper-0.2.14 → robocandywrapper-0.2.16}/robocandywrapper/factory.py

@@ -10,7 +10,7 @@ from lerobot.datasets.lerobot_dataset import (
     LeRobotDataset,
     LeRobotDatasetMetadata,
 )
-from lerobot.datasets.backward_compatibility import BackwardCompatibilityError
+from lerobot.datasets.utils import BackwardCompatibilityError
 from lerobot.datasets.transforms import ImageTransforms
 from lerobot.datasets.factory import IMAGENET_STATS
 from lerobot.utils.constants import ACTION, REWARD
@@ -71,7 +71,7 @@ def _create_datasets(
     root: Optional[str],
     revision: Optional[str],
     episodes: Optional[list[int] | dict[str, list[int]]],
-    video_backend: str,
+    video_backend: str | None,
     action_delta_indices: Optional[List] = None,
     observation_delta_indices: Optional[List] = None,
     reward_delta_indices: Optional[List] = None,
@@ -250,7 +250,7 @@ def make_dataset_without_config(
     action_delta_indices: List = None,
     observation_delta_indices: List = None,
     root: str = None,
-    video_backend: str = "pyav",
+    video_backend: str | None = None,
     episodes: list[int] | dict[str, list[int]] | None = None,
     revision: str | None = None,
     use_imagenet_stats: bool = True,

{robocandywrapper-0.2.14 → robocandywrapper-0.2.16}/robocandywrapper/plugins/__init__.py

@@ -12,6 +12,10 @@ from robocandywrapper.plugins.episode_outcome import (
     EpisodeOutcomePlugin,
     EpisodeOutcomeInstance,
 )
+from robocandywrapper.plugins.molmopoint import (
+    MolmoPointPlugin,
+    MolmoPointInstance,
+)
 
 __all__ = [
     "ControlModePlugin",
@@ -20,5 +24,7 @@ __all__ = [
     "AffordanceInstance",
     "EpisodeOutcomePlugin",
     "EpisodeOutcomeInstance",
+    "MolmoPointPlugin",
+    "MolmoPointInstance",
 ]
 

{robocandywrapper-0.2.14 → robocandywrapper-0.2.16}/robocandywrapper/plugins/control_mode.py

@@ -197,9 +197,12 @@ class ControlModeInstance(PluginInstance):
         """Find the episode_modes.json file for a dataset root, or None."""
         cw = dataset_root / CANDYWRAPPER_PLUGINS_DIR
         legacy = cw / "dagger_data_source" / "episode_modes.json"
+        legacy_2 = dataset_root / "dagger_data_source" / "episode_modes.json"
         current = cw / CONTROL_MODE_PLUGIN_NAME / "episode_modes.json"
         if legacy.exists():
             return legacy
+        if legacy_2.exists():
+            return legacy_2
         if current.exists():
             return current
         return None

robocandywrapper-0.2.16/robocandywrapper/plugins/molmopoint.py

@@ -0,0 +1,163 @@
+"""
+MolmoPoint Plugin for per-frame visual affordance point labels.
+
+Storage: {dataset_root}/candywrapper_plugins/molmopoint/episode_{idx:06d}.parquet
+
+Parquet schema per episode:
+    frame_index  : int     — episode-relative frame index (0-based)
+    camera       : str     — camera key, e.g. "observation.images.front"
+    point_x      : float   — normalised x coordinate (-1 to 1, left to right)
+    point_y      : float   — normalised y coordinate (-1 to 1, top to bottom)
+    is_keyframe  : bool    — True for raw MolmoPoint predictions, False for interpolated
+
+Coordinates are in [-1, 1] where (-1,-1) is top-left and (1,1) is bottom-right.
+"""
+from __future__ import annotations
+
+import warnings
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+import torch
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+
+from robocandywrapper.constants import CANDYWRAPPER_PLUGINS_DIR, MOLMOPOINT_PLUGIN_NAME
+from robocandywrapper.plugin import DatasetPlugin, PluginInstance
+
+
+def _calculate_episode_data_index(hf_dataset) -> dict[str, torch.Tensor]:
+    episode_data_index: dict[str, list[int]] = {"from": [], "to": []}
+    if len(hf_dataset) == 0:
+        return {"from": torch.tensor([]), "to": torch.tensor([])}
+    current_episode = None
+    for idx, episode_idx in enumerate(hf_dataset["episode_index"]):
+        if episode_idx != current_episode:
+            episode_data_index["from"].append(idx)
+            if current_episode is not None:
+                episode_data_index["to"].append(idx)
+            current_episode = episode_idx
+    episode_data_index["to"].append(idx + 1)
+    return {k: torch.tensor(v) for k, v in episode_data_index.items()}
+
+
+class MolmoPointPlugin(DatasetPlugin):
+    """Plugin that provides per-frame MolmoPoint affordance labels."""
+
+    def attach(self, dataset: "LeRobotDataset") -> "MolmoPointInstance":
+        return MolmoPointInstance(dataset, self)
+
+
+class MolmoPointInstance(PluginInstance):
+    """Dataset-specific MolmoPoint label data."""
+
+    def __init__(self, dataset: "LeRobotDataset", config: MolmoPointPlugin):
+        super().__init__(dataset)
+        self.config = config
+        self._data: Dict[int, "pd.DataFrame"] = {}
+        self._cached_episode_data_index: dict[str, torch.Tensor] | None = None
+
+        self._camera_keys: list[str] = list(dataset.meta.video_keys)
+        self._camera_to_idx: dict[str, int] = {
+            k: i for i, k in enumerate(self._camera_keys)
+        }
+        self._load()
+
+    def _get_episode_data_index(self) -> dict[str, torch.Tensor]:
+        if hasattr(self.dataset, "episode_data_index"):
+            return self.dataset.episode_data_index
+        if self._cached_episode_data_index is None:
+            self._cached_episode_data_index = _calculate_episode_data_index(
+                self.dataset.hf_dataset
+            )
+        return self._cached_episode_data_index
+
+    def get_data_keys(self) -> list[str]:
+        return ["pointing", "pointing_mask"]
+
+    def get_item_data(
+        self,
+        idx: int,
+        episode_idx: int,
+        accumulated_data: Optional[Dict[str, Any]] = None,
+    ) -> dict[str, Any]:
+        """Return per-camera pointing vector and mask for this frame.
+
+        Returns:
+            pointing: tensor of shape (num_cameras * 2,) with [x0, y0, x1, y1, ...]
+                in [-1, 1]. One (x, y) pair per camera in dataset.meta.video_keys
+                order. Unlabelled cameras get (0, 0).
+            pointing_mask: bool tensor of shape (num_cameras,), True where
+                a real label exists for that camera.
+        """
+        num_cameras = len(self._camera_keys)
+        pointing = torch.zeros(num_cameras * 2)
+        pointing_mask = torch.zeros(num_cameras, dtype=torch.bool)
+
+        if episode_idx not in self._data:
+            return {"pointing": pointing, "pointing_mask": pointing_mask}
+
+        ep_data_index = self._get_episode_data_index()
+        ep_start = ep_data_index["from"][episode_idx].item()
+        frame_in_episode = idx - ep_start
+
+        df = self._data[episode_idx]
+        rows = df[df["frame_index"] == frame_in_episode]
+
+        for _, row in rows.iterrows():
+            cam_idx = self._camera_to_idx.get(row["camera"])
+            if cam_idx is None:
+                continue
+            pointing[cam_idx * 2] = float(row["point_x"])
+            pointing[cam_idx * 2 + 1] = float(row["point_y"])
+            pointing_mask[cam_idx] = True
+
+        return {"pointing": pointing, "pointing_mask": pointing_mask}
+
+    def _get_plugin_dir(self) -> Path:
+        plugin_dir = (
+            Path(self.dataset.root) / CANDYWRAPPER_PLUGINS_DIR / MOLMOPOINT_PLUGIN_NAME
+        )
+        plugin_dir.mkdir(parents=True, exist_ok=True)
+        return plugin_dir
+
+    def _episode_path(self, episode_idx: int) -> Path:
+        return self._get_plugin_dir() / f"episode_{episode_idx:06d}.parquet"
+
+    def _load(self) -> None:
+        plugin_dir = self._get_plugin_dir()
+        import pandas as pd
+
+        loaded = 0
+        for pq_file in sorted(plugin_dir.glob("episode_*.parquet")):
+            try:
+                ep_idx = int(pq_file.stem.split("_")[1])
+                self._data[ep_idx] = pd.read_parquet(pq_file)
+                loaded += 1
+            except Exception as e:
+                warnings.warn(f"Could not load {pq_file}: {e}")
+
+        if loaded:
+            print(f"Loaded MolmoPoint labels for {loaded} episodes from {plugin_dir}")
+
+    def save_episode(self, episode_idx: int, df: "pd.DataFrame") -> Path:
+        """Save MolmoPoint labels for one episode. Returns the saved path."""
+        self._data[episode_idx] = df
+        path = self._episode_path(episode_idx)
+        df.to_parquet(path, index=False)
+        return path
+
+    def get_episode_labels(self, episode_idx: int) -> "pd.DataFrame | None":
+        return self._data.get(episode_idx)
+
+    def get_episode_labels_by_camera(
+        self, episode_idx: int, camera: str,
+    ) -> "pd.DataFrame | None":
+        df = self._data.get(episode_idx)
+        if df is None:
+            return None
+        filtered = df[df["camera"] == camera]
+        return filtered if not filtered.empty else None
+
+    @property
+    def labeled_episodes(self) -> list[int]:
+        return sorted(self._data.keys())

{robocandywrapper-0.2.14 → robocandywrapper-0.2.16}/robocandywrapper/wrapper.py

@@ -627,7 +627,7 @@ class WrappedRobotDataset(torch.utils.data.Dataset):
         else:
             item = dataset[local_idx]
 
-        episode_idx = item["episode_index"].item()
+        episode_idx = item["episode_index"].item() if isinstance(item["episode_index"], torch.Tensor) else item["episode_index"]
         
         # Add dataset index
         item["dataset_index"] = torch.tensor(dataset_idx)

{robocandywrapper-0.2.14 → robocandywrapper-0.2.16/robocandywrapper.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: robocandywrapper
-Version: 0.2.14
+Version: 0.2.16
 Summary: Sweet wrappers for extending and remixing LeRobot Datasets
 Author: RoboCandyWrapper Contributors
 License: MIT License
@@ -43,13 +43,14 @@ Requires-Dist: torch>=2.0.0
 Requires-Dist: lerobot<0.5,>=0.4
 Requires-Dist: pandas>=1.3.0
 Requires-Dist: rewact_tools
+Requires-Dist: motion-primitives
+Requires-Dist: lerobot-policy-diffusion-motion-primitives
+Requires-Dist: lerobot-policy-diffusion-pointing
 Provides-Extra: dev
 Requires-Dist: pytest>=7.0.0; extra == "dev"
 Requires-Dist: black>=22.0.0; extra == "dev"
 Requires-Dist: isort>=5.10.0; extra == "dev"
 Requires-Dist: flake8>=4.0.0; extra == "dev"
-Requires-Dist: motion-primitives; extra == "dev"
-Requires-Dist: lerobot-policy-diffusion-motion-primitives; extra == "dev"
 Dynamic: license-file
 Dynamic: requires-python
 

{robocandywrapper-0.2.14 → robocandywrapper-0.2.16}/robocandywrapper.egg-info/SOURCES.txt

@@ -24,6 +24,7 @@ robocandywrapper/plugins/__init__.py
 robocandywrapper/plugins/affordance.py
 robocandywrapper/plugins/control_mode.py
 robocandywrapper/plugins/episode_outcome.py
+robocandywrapper/plugins/molmopoint.py
 robocandywrapper/plugins/subtask.py
 robocandywrapper/samplers/__init__.py
 robocandywrapper/samplers/config.py

{robocandywrapper-0.2.14 → robocandywrapper-0.2.16}/robocandywrapper.egg-info/requires.txt

@@ -3,11 +3,12 @@ torch>=2.0.0
 lerobot<0.5,>=0.4
 pandas>=1.3.0
 rewact_tools
+motion-primitives
+lerobot-policy-diffusion-motion-primitives
+lerobot-policy-diffusion-pointing
 
 [dev]
 pytest>=7.0.0
 black>=22.0.0
 isort>=5.10.0
 flake8>=4.0.0
-motion-primitives
-lerobot-policy-diffusion-motion-primitives

{robocandywrapper-0.2.14 → robocandywrapper-0.2.16}/setup.py

@@ -9,7 +9,7 @@ long_description = readme_file.read_text(encoding="utf-8") if readme_file.exists
 
 setup(
     name="robocandywrapper",
-    version="0.2.14",
+    version="0.2.16",
     description="Sweet wrappers for extending and remixing LeRobot Datasets",
     long_description=long_description,
     long_description_content_type="text/markdown",