robocandywrapper 0.2.5__tar.gz → 0.2.7__tar.gz

{robocandywrapper-0.2.5/robocandywrapper.egg-info → robocandywrapper-0.2.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: robocandywrapper
-Version: 0.2.5
+Version: 0.2.7
 Summary: Sweet wrappers for extending and remixing LeRobot Datasets
 Author: RoboCandyWrapper Contributors
 License: MIT License

{robocandywrapper-0.2.5 → robocandywrapper-0.2.7}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "robocandywrapper"
-version = "0.2.5"
+version = "0.2.7"
 description = "Sweet wrappers for extending and remixing LeRobot Datasets"
 readme = "README.md"
 requires-python = ">=3.10"

{robocandywrapper-0.2.5 → robocandywrapper-0.2.7}/robocandywrapper/__init__.py

@@ -17,7 +17,7 @@ from robocandywrapper.constants import (
     EPISODE_OUTCOME_PLUGIN_NAME,
 )
 
-__version__ = "0.2.5"
+__version__ = "0.2.7"
 
 __all__ = [
     "DatasetPlugin",

{robocandywrapper-0.2.5 → robocandywrapper-0.2.7}/robocandywrapper/metadata_view.py

@@ -5,6 +5,39 @@ import logging
 from typing import Optional
 
 import numpy as np
+import pandas as pd
+
+
+def _task_names_from_meta_tasks(meta_tasks) -> set[str]:
+    """
+    Extract task name strings from a dataset's meta.tasks in either format.
+
+    - LeRobot 2.1 (dict): tasks is dict[task_index, task_name].
+    - LeRobot 3.0 (parquet/DataFrame): tasks is a pandas DataFrame where the
+      *row index* is the task name (string) and the only column is "task_index"
+      (integer). So: df.index = task names, df["task_index"] = indices.
+
+    Returns a set of unique task name strings.
+    """
+    if meta_tasks is None:
+        return set()
+    if isinstance(meta_tasks, dict):
+        return {str(v) for v in meta_tasks.values() if isinstance(v, str)}
+    # DataFrame or object with to_pandas() (e.g. HuggingFace Dataset)
+    df = None
+    if hasattr(meta_tasks, "to_pandas"):
+        df = meta_tasks.to_pandas()
+    elif isinstance(meta_tasks, pd.DataFrame):
+        df = meta_tasks
+    if df is None:
+        return set()
+    # 3.0 format: optional "task" column, else task names are the row index
+    if "task" in df.columns:
+        return set(df["task"].dropna().astype(str).unique())
+    if "task_index" in df.columns:
+        # Task names are the DataFrame index (row labels), not the column
+        return set(df.index.astype(str))
+    return set()
 
 
 def aggregate_stats_weighted(
@@ -241,7 +274,9 @@ class WrappedRobotDatasetMetadataView:
         Collated task index -> task name mapping across all datasets.
 
         Collects all unique task strings from every dataset that has a `tasks`
-        attribute on its meta (e.g. LeRobot v2.1 datasets). Assigns a unified
+        attribute on its meta. Supports both formats: LeRobot 2.1 (dict
+        task_index -> task name) and LeRobot 3.0 (parquet/pandas DataFrame with
+        "task" column, or object with to_pandas()). Assigns a unified
         task_index to each unique task name (sorted for determinism), so
         consumers get a single dict[int, str] suitable for _coerce_task_mapping
         and similar use.
@@ -252,10 +287,7 @@ class WrappedRobotDatasetMetadataView:
         all_task_names: set[str] = set()
         for dataset in self._datasets:
             meta_tasks = getattr(dataset.meta, "tasks", None)
-            if meta_tasks is not None and isinstance(meta_tasks, dict):
-                for _idx, name in meta_tasks.items():
-                    if isinstance(name, str):
-                        all_task_names.add(name)
+            all_task_names.update(_task_names_from_meta_tasks(meta_tasks))
 
         sorted_names = sorted(all_task_names)
         self._tasks = {idx: name for idx, name in enumerate(sorted_names)}

{robocandywrapper-0.2.5 → robocandywrapper-0.2.7}/robocandywrapper/wrapper.py

@@ -452,7 +452,7 @@ class WrappedRobotDataset(torch.utils.data.Dataset):
             return None
         
         for instance in self._plugin_instances[dataset_idx]:
-            if isinstance(instance, plugin_type.attach(self._datasets[0]).__class__):
+            if isinstance(instance, plugin_type.attach(self._datasets[dataset_idx]).__class__):
                 # Check if the instance came from the right plugin type
                 # This is a bit hacky - better to store plugin class reference
                 pass

{robocandywrapper-0.2.5 → robocandywrapper-0.2.7/robocandywrapper.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: robocandywrapper
-Version: 0.2.5
+Version: 0.2.7
 Summary: Sweet wrappers for extending and remixing LeRobot Datasets
 Author: RoboCandyWrapper Contributors
 License: MIT License

{robocandywrapper-0.2.5 → robocandywrapper-0.2.7}/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.5",
+    version="0.2.7",
     description="Sweet wrappers for extending and remixing LeRobot Datasets",
     long_description=long_description,
     long_description_content_type="text/markdown",