admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-05-16 17:20:05 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
3568df8a35989863c8bb3771c75bcb4f4cf5cb16
lerobot/examples/dataset/annotate.py
T
Jade Choghari 3568df8a35 woking on qwen
2025-12-08 14:03:47 +00:00

1153 lines
40 KiB
Python
Raw Blame History

#!/usr/bin/env python
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Automatic Skill Annotation for LeRobot Datasets.
This script performs automatic subtask/skill labeling for ANY LeRobot dataset using
Vision-Language Models (VLMs). It segments each robot demonstration into short atomic
skills (1-3 seconds each) and updates the dataset's task field.
The pipeline:
1. Loads a LeRobot dataset (local or from HuggingFace Hub)
2. For each episode, extracts video frames
3. Uses a VLM to identify skill boundaries and labels
4. Updates the dataset's task metadata with skill annotations
Supported VLMs (modular design allows easy extension):
- Qwen2-VL (default): "Qwen/Qwen2-VL-7B-Instruct"
- Qwen3-VL: "Qwen/Qwen3-VL-30B-A3B-Instruct"
- SmolVLM: "HuggingFaceTB/SmolVLM-Instruct"
Usage:
```bash
python examples/dataset/annotate.py \
--repo-id your-username/your-dataset \
--video-key observation.images.base \
--model Qwen/Qwen2-VL-7B-Instruct \
--push-to-hub
```
Or with a local dataset:
```bash
python examples/dataset/annotate.py \
--data-dir /path/to/local/dataset \
--video-key observation.images.base
```
After running, you can access the skill for any frame via:
```python
dataset = LeRobotDataset(repo_id="your/dataset")
item = dataset[100]
task_idx = item["task_index"]
skill_name = dataset.meta.tasks.iloc[task_idx].name
```
"""
import argparse
import json
import re
import subprocess
import tempfile
import textwrap
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Any
import cv2
import torch
from rich.console import Console
from rich.progress import Progress, SpinnerColumn, TextColumn
from lerobot.datasets.lerobot_dataset import LeRobotDataset
# =============================================================================
# Skill Annotation Data Structures
# =============================================================================
class Skill:
"""Represents a single atomic skill/subtask in a demonstration."""
def __init__(self, name: str, start: float, end: float):
self.name = name
self.start = start # Start timestamp in seconds
self.end = end # End timestamp in seconds
def to_dict(self) -> dict:
return {"name": self.name, "start": self.start, "end": self.end}
@classmethod
def from_dict(cls, data: dict) -> "Skill":
return cls(name=data["name"], start=data["start"], end=data["end"])
def __repr__(self) -> str:
return f"Skill(name='{self.name}', start={self.start:.2f}, end={self.end:.2f})"
class EpisodeSkills:
"""Container for all skills in an episode."""
def __init__(self, episode_index: int, description: str, skills: list[Skill]):
self.episode_index = episode_index
self.description = description
self.skills = skills
def to_dict(self) -> dict:
return {
"episode_index": self.episode_index,
"description": self.description,
"skills": [s.to_dict() for s in self.skills],
}
# =============================================================================
# VLM Interface (Abstract Base Class for Modularity)
# =============================================================================
class BaseVLM(ABC):
"""
Abstract base class for Vision-Language Models.
To add a new VLM:
1. Create a subclass of BaseVLM
2. Implement the `__init__` and `segment_skills` methods
3. Register it in the VLM_REGISTRY dictionary
"""
@abstractmethod
def __init__(self, model_name: str, device: str = "cuda", torch_dtype: torch.dtype = torch.bfloat16):
"""Initialize the VLM with model name, device, and dtype."""
pass
@abstractmethod
def segment_skills(
self, video_path: Path, episode_duration: float, coarse_goal: str | None = None
) -> list[Skill]:
"""
Segment a video into atomic skills.
Args:
video_path: Path to the video file
episode_duration: Total duration of the episode in seconds
coarse_goal: Optional high-level task description
Returns:
List of Skill objects representing atomic manipulation skills
"""
pass
def create_skill_segmentation_prompt(coarse_goal: str | None = None) -> str:
"""Create the prompt for skill segmentation."""
goal_context = f'The overall goal is: "{coarse_goal}"\n\n' if coarse_goal else ""
return textwrap.dedent(f"""\
# Role
You are a Robotics Vision System specializing in temporal action segmentation for robot manipulation demonstrations.
# Task
{goal_context}Segment this robot demonstration video into short atomic manipulation skills. Each skill should:
- Last approximately 1-3 seconds
- Describe a clear, single action (e.g., "pick up object", "move arm left", "release gripper")
- Have precise start and end timestamps
# Requirements
1. **Atomic Actions**: Each skill should be a single, indivisible action
2. **Complete Coverage**: Skills must cover the entire video duration with no gaps
3. **Boundary Consistency**: The end of one skill equals the start of the next
4. **Natural Language**: Use clear, descriptive names for each skill
5. **Timestamps**: Use seconds (float) for all timestamps
# Output Format
After your analysis, output ONLY valid JSON with this exact structure:
```json
{{
"skills": [
{{"name": "skill description", "start": 0.0, "end": 1.5}},
{{"name": "another skill", "start": 1.5, "end": 3.2}}
]
}}
```
The first skill must start at 0.0 and the last skill must end at the video duration.
""")
# =============================================================================
# Qwen2-VL Implementation
# =============================================================================
class Qwen2VL(BaseVLM):
"""Qwen2-VL model for skill segmentation."""
def __init__(self, model_name: str, device: str = "cuda", torch_dtype: torch.dtype = torch.bfloat16):
from qwen_vl_utils import process_vision_info
from transformers import AutoProcessor, Qwen2VLForConditionalGeneration
self.console = Console()
self.device = device
self.model_name = model_name
self.process_vision_info = process_vision_info
self.console.print(f"[cyan]Loading Qwen2-VL model: {model_name}...[/cyan]")
self.model = Qwen2VLForConditionalGeneration.from_pretrained(
model_name, torch_dtype=torch_dtype, device_map=device, trust_remote_code=True
)
self.processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
self.console.print(f"[green]✓ Model loaded successfully on {device}[/green]")
def segment_skills(
self, video_path: Path, episode_duration: float, coarse_goal: str | None = None
) -> list[Skill]:
"""Segment video into skills using Qwen2-VL."""
prompt = create_skill_segmentation_prompt(coarse_goal)
duration_str = f"{int(episode_duration // 60):02d}:{int(episode_duration % 60):02d}"
messages = [
{"role": "system", "content": [{"type": "text", "text": prompt}]},
{
"role": "user",
"content": [
{"type": "video", "video": str(video_path), "fps": 1.0},
{
"type": "text",
"text": f"Video duration: {duration_str} (~{episode_duration:.1f}s). Segment into atomic skills.",
},
],
},
]
text = self.processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
image_inputs, video_inputs = self.process_vision_info(messages)
inputs = self.processor(
text=[text],
images=image_inputs,
videos=video_inputs,
padding=True,
return_tensors="pt",
).to(self.device)
with torch.no_grad():
generated_ids = self.model.generate(**inputs, max_new_tokens=1024, do_sample=True, temperature=0.7)
response = self.processor.batch_decode(
[out[len(inp) :] for inp, out in zip(inputs.input_ids, generated_ids)],
skip_special_tokens=True,
)[0].strip()
return self._parse_skills_response(response)
def _parse_skills_response(self, response: str) -> list[Skill]:
"""Parse the VLM response into Skill objects."""
# Extract JSON from response
if "```json" in response:
response = response.split("```json")[1].split("```")[0]
elif "```" in response:
response = response.split("```")[1].split("```")[0]
try:
data = json.loads(response)
skills_data = data.get("skills", data)
if isinstance(skills_data, list):
return [Skill.from_dict(s) for s in skills_data]
except json.JSONDecodeError:
# Try to find JSON object in response
match = re.search(r"\{.*\}", response, re.DOTALL)
if match:
data = json.loads(match.group())
skills_data = data.get("skills", [])
return [Skill.from_dict(s) for s in skills_data]
raise ValueError(f"Could not parse skills from response: {response[:200]}...")
# =============================================================================
# Qwen3-VL Implementation (MoE variant)
# =============================================================================
class Qwen3VL(BaseVLM):
"""Qwen3-VL MoE model for skill segmentation."""
def __init__(self, model_name: str, device: str = "cuda", torch_dtype: torch.dtype = torch.bfloat16):
from qwen_vl_utils import process_vision_info
from transformers import AutoProcessor, Qwen3VLMoeForConditionalGeneration
self.console = Console()
self.device = device
self.model_name = model_name
self.process_vision_info = process_vision_info
self.console.print(f"[cyan]Loading Qwen3-VL model: {model_name}...[/cyan]")
self.model = Qwen3VLMoeForConditionalGeneration.from_pretrained(
model_name, torch_dtype=torch_dtype, device_map=device, trust_remote_code=True
)
self.processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
self.console.print(f"[green]✓ Model loaded successfully on {device}[/green]")
def segment_skills(
self, video_path: Path, episode_duration: float, coarse_goal: str | None = None
) -> list[Skill]:
"""Segment video into skills using Qwen3-VL."""
prompt = create_skill_segmentation_prompt(coarse_goal)
duration_str = f"{int(episode_duration // 60):02d}:{int(episode_duration % 60):02d}"
messages = [
{"role": "system", "content": [{"type": "text", "text": prompt}]},
{
"role": "user",
"content": [
{"type": "video", "video": str(video_path), "fps": 1.0},
{
"type": "text",
"text": f"Video duration: {duration_str} (~{episode_duration:.1f}s). Segment into atomic skills.",
},
],
},
]
text = self.processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
image_inputs, video_inputs = self.process_vision_info(messages)
inputs = self.processor(
text=[text],
images=image_inputs,
videos=video_inputs,
padding=True,
return_tensors="pt",
).to(self.device)
with torch.no_grad():
generated_ids = self.model.generate(**inputs, max_new_tokens=1024, do_sample=True, temperature=0.7)
response = self.processor.batch_decode(
[out[len(inp) :] for inp, out in zip(inputs.input_ids, generated_ids)],
skip_special_tokens=True,
)[0].strip()
return self._parse_skills_response(response)
def _parse_skills_response(self, response: str) -> list[Skill]:
"""Parse the VLM response into Skill objects."""
if "```json" in response:
response = response.split("```json")[1].split("```")[0]
elif "```" in response:
response = response.split("```")[1].split("```")[0]
try:
data = json.loads(response)
skills_data = data.get("skills", data)
if isinstance(skills_data, list):
return [Skill.from_dict(s) for s in skills_data]
except json.JSONDecodeError:
match = re.search(r"\{.*\}", response, re.DOTALL)
if match:
data = json.loads(match.group())
skills_data = data.get("skills", [])
return [Skill.from_dict(s) for s in skills_data]
raise ValueError(f"Could not parse skills from response: {response[:200]}...")
# =============================================================================
# SmolVLM Implementation
# =============================================================================
class SmolVLM(BaseVLM):
"""SmolVLM model for skill segmentation (lighter weight alternative)."""
def __init__(self, model_name: str, device: str = "cuda", torch_dtype: torch.dtype = torch.bfloat16):
from transformers import AutoModelForVision2Seq, AutoProcessor
self.console = Console()
self.device = device
self.model_name = model_name
self.console.print(f"[cyan]Loading SmolVLM model: {model_name}...[/cyan]")
self.processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
self.model = AutoModelForVision2Seq.from_pretrained(
model_name,
torch_dtype=torch_dtype,
# _attn_implementation="flash_attention_2" if device == "cuda" else "eager",
).to(device)
self.console.print(f"[green]✓ Model loaded successfully on {device}[/green]")
def segment_skills(
self, video_path: Path, episode_duration: float, coarse_goal: str | None = None
) -> list[Skill]:
"""Segment video into skills using SmolVLM with frame sampling."""
import PIL.Image
# SmolVLM works with images, so we sample frames from the video
frames = self._extract_frames(video_path, target_fps=1)
if not frames:
raise ValueError(f"Could not extract frames from {video_path}")
prompt = create_skill_segmentation_prompt(coarse_goal)
duration_str = f"{int(episode_duration // 60):02d}:{int(episode_duration % 60):02d}"
# Sample frames (up to 8 frames to avoid context overflow)
frame_indices = self._select_frame_indices(len(frames), max_frames=8)
# Convert frames to PIL images
pil_images = [
PIL.Image.fromarray(cv2.cvtColor(frames[idx], cv2.COLOR_BGR2RGB))
for idx in frame_indices
]
# Create message content with image placeholders
content = [{"type": "text", "text": prompt}]
# Add image placeholders (one for each frame)
for _ in frame_indices:
content.append({"type": "image"})
content.append(
{
"type": "text",
"text": f"These are {len(frame_indices)} sampled frames from a {duration_str} video. Segment into atomic skills.",
}
)
messages = [{"role": "user", "content": content}]
# Apply chat template to get the prompt
prompt = self.processor.apply_chat_template(messages, add_generation_prompt=True)
# Process inputs with both text and images
inputs = self.processor(text=prompt, images=pil_images, return_tensors="pt")
inputs = inputs.to(self.device)
with torch.no_grad():
generated_ids = self.model.generate(**inputs, max_new_tokens=1024, do_sample=True, temperature=0.7)
response = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
return self._parse_skills_response(response, episode_duration)
def _extract_frames(self, video_path: Path, target_fps: int = 1) -> list:
"""Extract frames from video at target FPS."""
cap = cv2.VideoCapture(str(video_path))
frames = []
fps = cap.get(cv2.CAP_PROP_FPS) or 30
frame_interval = int(fps / target_fps)
frame_count = 0
while True:
ret, frame = cap.read()
if not ret:
break
if frame_count % frame_interval == 0:
frames.append(frame)
frame_count += 1
cap.release()
return frames
def _select_frame_indices(self, total_frames: int, max_frames: int = 8) -> list[int]:
"""Select evenly spaced frame indices."""
if total_frames <= max_frames:
return list(range(total_frames))
step = total_frames / max_frames
return [int(i * step) for i in range(max_frames)]
def _parse_skills_response(self, response: str, episode_duration: float) -> list[Skill]:
"""Parse the VLM response into Skill objects."""
if "```json" in response:
response = response.split("```json")[1].split("```")[0]
elif "```" in response:
response = response.split("```")[1].split("```")[0]
try:
data = json.loads(response)
skills_data = data.get("skills", data)
breakpoint()
if isinstance(skills_data, list):
return [Skill.from_dict(s) for s in skills_data]
except json.JSONDecodeError:
match = re.search(r"\{.*\}", response, re.DOTALL)
if match:
data = json.loads(match.group())
skills_data = data.get("skills", [])
return [Skill.from_dict(s) for s in skills_data]
# Fallback: create a single skill covering the whole episode
self.console.print("[yellow]Warning: Could not parse skills, creating single skill[/yellow]")
return [Skill(name="perform manipulation", start=0.0, end=episode_duration)]
# =============================================================================
# VLM Registry - Add new VLMs here
# =============================================================================
VLM_REGISTRY: dict[str, type[BaseVLM]] = {
# Qwen2-VL variants
"Qwen/Qwen2-VL-2B-Instruct": Qwen2VL,
"Qwen/Qwen2-VL-7B-Instruct": Qwen2VL,
"Qwen/Qwen2-VL-72B-Instruct": Qwen2VL,
# Qwen3-VL variants (MoE)
"Qwen/Qwen3-VL-30B-A3B-Instruct": Qwen3VL,
# SmolVLM variants
"HuggingFaceTB/SmolVLM-Instruct": SmolVLM,
"HuggingFaceTB/SmolVLM-256M-Instruct": SmolVLM,
"HuggingFaceTB/SmolVLM-500M-Instruct": SmolVLM,
}
def get_vlm(model_name: str, device: str = "cuda", torch_dtype: torch.dtype = torch.bfloat16) -> BaseVLM:
"""
Factory function to get the appropriate VLM based on model name.
Args:
model_name: HuggingFace model identifier
device: Device to load model on
torch_dtype: Data type for model weights
Returns:
Initialized VLM instance
Raises:
ValueError: If model is not in registry
"""
# Check exact match first
if model_name in VLM_REGISTRY:
return VLM_REGISTRY[model_name](model_name, device, torch_dtype)
# Check for partial matches (e.g., "qwen2" in model name)
model_lower = model_name.lower()
if "qwen3" in model_lower:
return Qwen3VL(model_name, device, torch_dtype)
elif "qwen2" in model_lower or "qwen-vl" in model_lower:
return Qwen2VL(model_name, device, torch_dtype)
elif "smolvlm" in model_lower:
return SmolVLM(model_name, device, torch_dtype)
raise ValueError(
f"Unknown model: {model_name}. "
f"Supported models: {list(VLM_REGISTRY.keys())}. "
"Or implement a new VLM class inheriting from BaseVLM."
)
# =============================================================================
# Video Extraction Utilities
# =============================================================================
class VideoExtractor:
"""Utilities for extracting and processing video segments from LeRobot datasets."""
def __init__(self, console: Console | None = None):
self.console = console or Console()
def extract_episode_video(
self,
video_path: Path,
start_timestamp: float,
end_timestamp: float,
target_fps: int = 1,
) -> Path:
"""
Extract a specific episode segment from a concatenated video file.
Args:
video_path: Path to the source video file
start_timestamp: Start time in seconds
end_timestamp: End time in seconds
target_fps: Target frames per second for output
Returns:
Path to the extracted temporary video file
"""
tmp_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
tmp_path = Path(tmp_file.name)
tmp_file.close()
duration = end_timestamp - start_timestamp
self.console.print(
f"[cyan]Extracting: {start_timestamp:.1f}s - {end_timestamp:.1f}s ({duration:.1f}s)[/cyan]"
)
cmd = [
"ffmpeg",
"-i",
str(video_path),
"-ss",
str(start_timestamp),
"-t",
str(duration),
"-r",
str(target_fps),
"-c:v",
"libx264",
"-preset",
"ultrafast",
"-crf",
"23",
"-an",
"-y",
str(tmp_path),
]
try:
subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True)
except subprocess.CalledProcessError as e:
raise RuntimeError(f"FFmpeg failed: {e}") from e
except FileNotFoundError:
raise RuntimeError("FFmpeg not found. Please install ffmpeg.")
if not tmp_path.exists() or tmp_path.stat().st_size < 1024:
if tmp_path.exists():
tmp_path.unlink()
raise RuntimeError("Video extraction produced invalid file")
return tmp_path
def get_video_duration(self, video_path: Path) -> float:
"""Get duration of a video file in seconds."""
cap = cv2.VideoCapture(str(video_path))
fps = cap.get(cv2.CAP_PROP_FPS) or 30
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
cap.release()
return frame_count / fps
# =============================================================================
# Skill Annotation Pipeline
# =============================================================================
class SkillAnnotator:
"""
Main class for annotating LeRobot datasets with skill labels.
This class orchestrates the full annotation pipeline:
1. Load dataset
2. Extract video segments for each episode
3. Run VLM-based skill segmentation
4. Update dataset task metadata
"""
def __init__(
self,
vlm: BaseVLM,
video_extractor: VideoExtractor | None = None,
console: Console | None = None,
):
self.vlm = vlm
self.console = console or Console()
self.video_extractor = video_extractor or VideoExtractor(self.console)
def annotate_dataset(
self,
dataset: LeRobotDataset,
video_key: str,
episodes: list[int] | None = None,
skip_existing: bool = False,
) -> dict[int, EpisodeSkills]:
"""
Annotate all episodes in a dataset with skill labels.
Args:
dataset: LeRobot dataset to annotate
video_key: Key for video observations (e.g., "observation.images.base")
episodes: Specific episode indices to annotate (None = all)
skip_existing: Skip episodes that already have skill annotations
Returns:
Dictionary mapping episode index to EpisodeSkills
"""
episode_indices = episodes or list(range(dataset.meta.total_episodes))
annotations: dict[int, EpisodeSkills] = {}
# Get coarse task description if available
coarse_goal = self._get_coarse_goal(dataset)
# with Progress(
# SpinnerColumn(),
# TextColumn("[progress.description]{task.description}"),
# console=self.console,
# ) as progress:
# task = progress.add_task(f"Annotating {len(episode_indices)} episodes...", total=len(episode_indices))
print(f"Annotating {len(episode_indices)} episodes...")
for ep_idx in episode_indices:
# progress.update(task, description=f"Processing episode {ep_idx}...")
print(f"Processing episode {ep_idx}...")
try:
skills = self._annotate_episode(dataset, ep_idx, video_key, coarse_goal)
annotations[ep_idx] = EpisodeSkills(
episode_index=ep_idx,
description=coarse_goal,
skills=skills,
)
self.console.print(
f"[green]✓ Episode {ep_idx}: {len(skills)} skills identified[/green]"
)
except Exception as e:
self.console.print(f"[red]✗ Episode {ep_idx} failed: {e}[/red]")
# progress.advance(task)
return annotations
def _get_coarse_goal(self, dataset: LeRobotDataset) -> str:
"""Extract or generate the coarse task description."""
# Try to get from existing task metadata
if dataset.meta.tasks is not None and len(dataset.meta.tasks) > 0:
# Get the first task description
first_task = dataset.meta.tasks.index[0]
if first_task:
return str(first_task)
return "Perform the demonstrated manipulation task."
def _annotate_episode(
self,
dataset: LeRobotDataset,
episode_index: int,
video_key: str,
coarse_goal: str,
) -> list[Skill]:
"""Annotate a single episode with skill labels."""
# Get video path and timestamps for this episode
video_path = dataset.root / dataset.meta.get_video_file_path(episode_index, video_key)
if not video_path.exists():
raise FileNotFoundError(f"Video not found: {video_path}")
# Get episode timestamps from metadata
ep = dataset.meta.episodes[episode_index]
start_ts = float(ep[f"videos/{video_key}/from_timestamp"])
end_ts = float(ep[f"videos/{video_key}/to_timestamp"])
duration = end_ts - start_ts
# Extract episode segment to temporary file
extracted_path = self.video_extractor.extract_episode_video(
video_path, start_ts, end_ts, target_fps=1
)
try:
# Run VLM skill segmentation
skills = self.vlm.segment_skills(extracted_path, duration, coarse_goal)
return skills
finally:
# Clean up temporary file
if extracted_path.exists():
extracted_path.unlink()
# =============================================================================
# Metadata Writer - Updates per-frame task_index based on skills
# =============================================================================
def get_skill_for_timestamp(skills: list[Skill], timestamp: float) -> Skill | None:
"""
Find which skill covers a given timestamp.
Args:
skills: List of skills with start/end times
timestamp: Frame timestamp in seconds
Returns:
The Skill that covers this timestamp, or None if not found
"""
for skill in skills:
if skill.start <= timestamp < skill.end:
return skill
# Handle the last frame (end boundary)
if timestamp >= skill.end and skill == skills[-1]:
return skill
return skills[-1] if skills else None # Fallback to last skill
def update_dataset_tasks(
dataset: LeRobotDataset,
annotations: dict[int, EpisodeSkills],
) -> dict[str, int]:
"""
Register all unique skill names as new tasks in the dataset.
Args:
dataset: The LeRobot dataset to update
annotations: Dictionary of episode skills
Returns:
Dictionary mapping skill name to task_index
"""
import pandas as pd
from lerobot.datasets.utils import write_tasks
console = Console()
# Collect all unique skill names
all_skill_names: set[str] = set()
for episode_skills in annotations.values():
for skill in episode_skills.skills:
all_skill_names.add(skill.name)
console.print(f"[cyan]Found {len(all_skill_names)} unique skills[/cyan]")
# Build new tasks DataFrame
# Start with existing tasks if any
if dataset.meta.tasks is not None and len(dataset.meta.tasks) > 0:
existing_tasks = set(dataset.meta.tasks.index.tolist())
max_task_idx = dataset.meta.tasks["task_index"].max()
else:
existing_tasks = set()
max_task_idx = -1
# Add new skills as tasks
new_tasks = all_skill_names - existing_tasks
if new_tasks:
new_task_data = []
for i, skill_name in enumerate(sorted(new_tasks)):
new_task_data.append({
"task": skill_name,
"task_index": max_task_idx + 1 + i,
})
new_tasks_df = pd.DataFrame(new_task_data).set_index("task")
if dataset.meta.tasks is not None and len(dataset.meta.tasks) > 0:
dataset.meta.tasks = pd.concat([dataset.meta.tasks, new_tasks_df])
else:
dataset.meta.tasks = new_tasks_df
# Write updated tasks to disk
write_tasks(dataset.meta.tasks, dataset.root)
console.print(f"[green]✓ Added {len(new_tasks)} new tasks to tasks.parquet[/green]")
# Build skill name to task_index mapping
skill_to_task_idx = {
task_name: int(dataset.meta.tasks.loc[task_name, "task_index"])
for task_name in all_skill_names
}
return skill_to_task_idx
def update_frame_task_indices(
dataset: LeRobotDataset,
annotations: dict[int, EpisodeSkills],
skill_to_task_idx: dict[str, int],
) -> None:
"""
Update the task_index for each frame based on skill annotations.
This reads the data parquet files, updates task_index based on which
skill covers each frame's timestamp, and writes back to disk.
Args:
dataset: The LeRobot dataset to update
annotations: Dictionary of episode skills
skill_to_task_idx: Mapping from skill name to task_index
"""
import pandas as pd
console = Console()
# Group episodes by their data file (chunk_index, file_index)
episodes_by_file: dict[tuple[int, int], list[int]] = {}
for ep_idx in annotations.keys():
ep = dataset.meta.episodes[ep_idx]
chunk_idx = ep["data/chunk_index"]
file_idx = ep["data/file_index"]
key = (chunk_idx, file_idx)
if key not in episodes_by_file:
episodes_by_file[key] = []
episodes_by_file[key].append(ep_idx)
# Process each data file
for (chunk_idx, file_idx), episode_indices in episodes_by_file.items():
data_path = dataset.root / dataset.meta.data_path.format(
chunk_index=chunk_idx, file_index=file_idx
)
if not data_path.exists():
console.print(f"[yellow]Warning: Data file not found: {data_path}[/yellow]")
continue
# Read the parquet file
df = pd.read_parquet(data_path)
original_task_indices = df["task_index"].copy()
updated_count = 0
# Update task_index for each episode in this file
for ep_idx in episode_indices:
if ep_idx not in annotations:
continue
episode_skills = annotations[ep_idx]
skills = episode_skills.skills
# Get episode frame range
ep = dataset.meta.episodes[ep_idx]
ep_from = ep["dataset_from_index"]
ep_to = ep["dataset_to_index"]
# Filter to rows for this episode
episode_mask = (df["index"] >= ep_from) & (df["index"] < ep_to)
episode_rows = df.loc[episode_mask]
# Update task_index for each frame based on its timestamp
for idx, row in episode_rows.iterrows():
timestamp = row["timestamp"]
skill = get_skill_for_timestamp(skills, timestamp)
if skill and skill.name in skill_to_task_idx:
new_task_idx = skill_to_task_idx[skill.name]
if df.at[idx, "task_index"] != new_task_idx:
df.at[idx, "task_index"] = new_task_idx
updated_count += 1
# Write back if any changes were made
if updated_count > 0:
df.to_parquet(data_path, engine="pyarrow", compression="snappy", index=False)
console.print(
f"[green]✓ Updated {updated_count} frame task_indices in {data_path.name}[/green]"
)
def save_skill_annotations(
dataset: LeRobotDataset,
annotations: dict[int, EpisodeSkills],
output_path: Path | None = None,
) -> None:
"""
Save skill annotations to the dataset, updating both:
1. The tasks.parquet with new skill names
2. The per-frame task_index in data parquet files
This function updates the task field for each frame based on
which skill covers that frame's timestamp.
Args:
dataset: The LeRobot dataset to update
annotations: Dictionary of episode skills
output_path: Optional custom output path for the annotations JSON
"""
console = Console()
if not annotations:
console.print("[yellow]No annotations to save[/yellow]")
return
# Step 1: Register all unique skills as tasks
console.print("[cyan]Registering skills as tasks...[/cyan]")
skill_to_task_idx = update_dataset_tasks(dataset, annotations)
# Step 2: Update per-frame task_index in data parquet files
console.print("[cyan]Updating per-frame task indices...[/cyan]")
update_frame_task_indices(dataset, annotations, skill_to_task_idx)
# Step 3: Also save the raw skill annotations as JSON for reference
skills_data = {
"coarse_description": annotations[next(iter(annotations))].description,
"skill_to_task_index": skill_to_task_idx,
"episodes": {str(ep_idx): ann.to_dict() for ep_idx, ann in annotations.items()},
}
skills_path = output_path or (dataset.root / "meta" / "skills.json")
skills_path.parent.mkdir(parents=True, exist_ok=True)
with open(skills_path, "w") as f:
json.dump(skills_data, f, indent=2)
console.print(f"[green]✓ Saved skill annotations to {skills_path}[/green]")
# Reload the dataset's hf_dataset to reflect changes
dataset._lazy_loading = True
def load_skill_annotations(dataset_root: Path) -> dict | None:
"""Load existing skill annotations from a dataset."""
skills_path = dataset_root / "meta" / "skills.json"
if skills_path.exists():
with open(skills_path) as f:
return json.load(f)
return None
# =============================================================================
# Main Entry Point
# =============================================================================
def main():
"""Main entry point for the skill annotation script."""
parser = argparse.ArgumentParser(
description="Automatic skill annotation for LeRobot datasets using VLMs",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog=textwrap.dedent("""\
Examples:
# Annotate a HuggingFace Hub dataset
python annotate.py --repo-id user/dataset --video-key observation.images.base
# Annotate a local dataset
python annotate.py --data-dir /path/to/dataset --video-key observation.images.base
# Use a specific model
python annotate.py --repo-id user/dataset --video-key observation.images.base \\
--model Qwen/Qwen2-VL-7B-Instruct
# Push annotated dataset to Hub
python annotate.py --repo-id user/dataset --video-key observation.images.base --push-to-hub
"""),
)
# Data source (mutually exclusive)
data_group = parser.add_mutually_exclusive_group(required=True)
data_group.add_argument("--data-dir", type=str, help="Path to local LeRobot dataset")
data_group.add_argument("--repo-id", type=str, help="HuggingFace Hub dataset repository ID")
# Required arguments
parser.add_argument(
"--video-key",
type=str,
required=True,
help="Video observation key (e.g., 'observation.images.base')",
)
# Model configuration
parser.add_argument(
"--model",
type=str,
default="Qwen/Qwen2-VL-7B-Instruct",
help="VLM model to use for skill segmentation (default: Qwen/Qwen2-VL-7B-Instruct)",
)
parser.add_argument(
"--device",
type=str,
default="cuda",
help="Device to run model on (default: cuda)",
)
parser.add_argument(
"--dtype",
type=str,
default="bfloat16",
choices=["bfloat16", "float16", "float32"],
help="Model dtype (default: bfloat16)",
)
# Episode selection
parser.add_argument(
"--episodes",
type=int,
nargs="+",
help="Specific episode indices to annotate (default: all)",
)
parser.add_argument(
"--skip-existing",
action="store_true",
help="Skip episodes that already have annotations",
)
# Output options
parser.add_argument(
"--push-to-hub",
action="store_true",
help="Push annotated dataset to HuggingFace Hub",
)
parser.add_argument(
"--output-path",
type=str,
help="Custom output path for annotations JSON",
)
args = parser.parse_args()
console = Console()
# Validate arguments
dtype_map = {
"bfloat16": torch.bfloat16,
"float16": torch.float16,
"float32": torch.float32,
}
torch_dtype = dtype_map[args.dtype]
# Load dataset
console.print("[cyan]Loading dataset...[/cyan]")
if args.data_dir:
dataset = LeRobotDataset(repo_id="local/dataset", root=args.data_dir, download_videos=False)
else:
dataset = LeRobotDataset(repo_id=args.repo_id, download_videos=True)
console.print(f"[green]✓ Loaded dataset with {dataset.meta.total_episodes} episodes[/green]")
# Validate video key
if args.video_key not in dataset.meta.video_keys:
available = ", ".join(dataset.meta.video_keys)
console.print(f"[red]Error: Video key '{args.video_key}' not found. Available: {available}[/red]")
return
# Initialize VLM
console.print(f"[cyan]Initializing VLM: {args.model}...[/cyan]")
vlm = get_vlm(args.model, args.device, torch_dtype)
# Create annotator and run annotation
annotator = SkillAnnotator(vlm=vlm, console=console)
annotations = annotator.annotate_dataset(
dataset=dataset,
video_key=args.video_key,
episodes=args.episodes,
skip_existing=args.skip_existing,
)
# Save annotations
output_path = Path(args.output_path) if args.output_path else None
save_skill_annotations(dataset, annotations, output_path)
# Summary
total_skills = sum(len(ann.skills) for ann in annotations.values())
console.print(f"\n[bold green]✓ Annotation complete![/bold green]")
console.print(f" Episodes annotated: {len(annotations)}")
console.print(f" Total skills identified: {total_skills}")
# Push to hub if requested
if args.push_to_hub:
if args.data_dir:
console.print("[yellow]Warning: --push-to-hub requires --repo-id, skipping...[/yellow]")
else:
console.print("[cyan]Pushing to HuggingFace Hub...[/cyan]")
try:
dataset.push_to_hub(push_videos=False)
console.print(f"[green]✓ Pushed to {args.repo_id}[/green]")
except Exception as e:
console.print(f"[red]Push failed: {e}[/red]")
if __name__ == "__main__":
main()
Reference in New Issue View Git Blame Copy Permalink