Add files

app.py

@@ -1,7 +1,291 @@
+import os
+import sys
+import threading
+import queue
+import time
+import tempfile
+import shutil
+from contextlib import redirect_stdout, redirect_stderr
+from typing import List, Optional
+
 import gradio as gr
+from huggingface_hub import HfApi
+
+from delete_episodes import (
+    download_dataset,
+    list_episodes,
+    delete_episodes_and_repair,
+    upload_dataset,
+)
+
+
+class _StreamToQueue:
+    def __init__(self, q: "queue.Queue[str]"):
+        self.q = q
+        self._buffer = ""
+
+    def write(self, s: str):
+        if not isinstance(s, str):
+            s = str(s)
+        self._buffer += s
+        while "\n" in self._buffer:
+            line, self._buffer = self._buffer.split("\n", 1)
+            self.q.put(line + "\n")
+
+    def flush(self):
+        if self._buffer:
+            self.q.put(self._buffer)
+            self._buffer = ""
+
+
+def search_datasets_fn(query: str) -> List[str]:
+    """Search for datasets on HuggingFace"""
+    api = HfApi()
+    try:
+        items = api.list_datasets(search=(query or "").strip() or None)
+        repo_ids = [getattr(d, "id", None) or getattr(d, "repo_id", None) for d in items]
+        repo_ids = [r for r in repo_ids if r]
+        # Remove duplicates while preserving order
+        seen = set()
+        unique = []
+        for r in repo_ids:
+            if r not in seen:
+                unique.append(r)
+                seen.add(r)
+        return unique[:500]
+    except Exception as e:
+        print(f"Error searching datasets: {e}")
+        return []
+
+
+def download_and_list_episodes(repo_id: str, progress=gr.Progress()):
+    """Download dataset and list available episodes"""
+    if not repo_id:
+        return "Please provide a dataset repo ID.", []
+    
+    token = os.environ.get("HF_TOKEN")
+    temp_dir = tempfile.mkdtemp(prefix="episode_delete_")
+    
+    try:
+        progress(0.3, desc="Downloading dataset...")
+        download_dataset(repo_id, temp_dir, hf_token=token)
+        
+        progress(0.7, desc="Listing episodes...")
+        episodes = list_episodes(temp_dir)
+        
+        if not episodes:
+            return f"No episodes found in {repo_id}", []
+        
+        # Format episodes as choices
+        episode_choices = [f"Episode {ep}" for ep in episodes]
+        
+        return (
+            f"Downloaded {repo_id}. Found {len(episodes)} episodes.",
+            gr.update(choices=episode_choices, value=[]),
+        )
+    except Exception as e:
+        return f"Error: {str(e)}", gr.update(choices=[], value=[])
+
+
+def delete_episodes_stream(repo_id: str, selected_episodes: List[str], dest_repo_id: str):
+    """Delete selected episodes and upload to destination repo"""
+    if not repo_id:
+        yield "Please provide a source dataset repo ID."
+        return
+    
+    if not selected_episodes:
+        yield "Please select at least one episode to delete."
+        return
+    
+    if not dest_repo_id:
+        yield "Please provide a destination repo ID."
+        return
+    
+    # Parse episode numbers from selection (format: "Episode 0", "Episode 1", etc.)
+    episode_indexes = []
+    for ep_str in selected_episodes:
+        try:
+            ep_num = int(ep_str.replace("Episode ", ""))
+            episode_indexes.append(ep_num)
+        except ValueError:
+            yield f"Invalid episode format: {ep_str}"
+            return
+    
+    token = os.environ.get("HF_TOKEN")
+    q: "queue.Queue[str]" = queue.Queue()
+    done = {"ok": False, "msg": ""}
+    
+    def _worker():
+        stream = _StreamToQueue(q)
+        temp_dir = tempfile.mkdtemp(prefix="episode_delete_")
+        
+        try:
+            with redirect_stdout(stream), redirect_stderr(stream):
+                print("Downloading dataset...", flush=True)
+                download_dataset(repo_id, temp_dir, hf_token=token)
+                
+                print(f"\nDeleting episodes: {episode_indexes}", flush=True)
+                delete_episodes_and_repair(
+                    dataset_path=temp_dir,
+                    episode_indexes=episode_indexes,
+                    run_stats=False,  # Skip stats for now as script may not be available
+                )
+                
+                print(f"\nUploading to {dest_repo_id}...", flush=True)
+                upload_dataset(
+                    local_dir=temp_dir,
+                    dest_repo_id=dest_repo_id,
+                    hf_token=token,
+                    commit_message=f"Deleted episodes: {episode_indexes}",
+                    private=False,
+                )
+                
+                print("\nUpload complete!", flush=True)
+                done["ok"] = True
+                done["msg"] = f"Successfully deleted {len(episode_indexes)} episodes and uploaded to {dest_repo_id}"
+        except Exception as e:
+            print(f"\nError: {e}", flush=True)
+            done["ok"] = False
+            done["msg"] = f"Error: {e}"
+        finally:
+            # Cleanup
+            try:
+                if os.path.isdir(temp_dir):
+                    shutil.rmtree(temp_dir, ignore_errors=True)
+                    print(f"\nCleaned up temp directory: {temp_dir}", flush=True)
+            except Exception:
+                pass
+            try:
+                stream.flush()
+            except Exception:
+                pass
+    
+    t = threading.Thread(target=_worker, daemon=True)
+    t.start()
+    
+    buffer = ""
+    yield "Starting process...\n"
+    
+    while t.is_alive() or not q.empty():
+        try:
+            line = q.get(timeout=0.1)
+            buffer += line
+            if len(buffer) > 0:
+                yield buffer
+        except queue.Empty:
+            pass
+        time.sleep(0.05)
+    
+    # Final status
+    if done["msg"]:
+        buffer += ("\n" if not buffer.endswith("\n") else "") + "=" * 50 + "\n" + done["msg"]
+    yield buffer
+
+
+# Build the Gradio interface
+with gr.Blocks(title="LeRobot Episode Deleter", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🗑️ LeRobot Dataset Episode Deleter")
+    gr.Markdown(
+        "Delete specific episodes from a HuggingFace LeRobot dataset and upload the cleaned version."
+    )
+    
+    with gr.Tabs():
+        with gr.Tab("Step 1: Select Dataset"):
+            with gr.Row():
+                with gr.Column(scale=3):
+                    search_box = gr.Textbox(
+                        label="Search Datasets",
+                        placeholder="Enter keyword or organization name (e.g., 'lerobot', 'griffinlabs-cortex')",
+                        value="griffinlabs-cortex"
+                    )
+                with gr.Column(scale=1):
+                    search_btn = gr.Button("🔍 Search", variant="primary")
+            
+            dataset_dropdown = gr.Dropdown(
+                label="Available Datasets",
+                choices=search_datasets_fn("griffinlabs-cortex"),
+                interactive=True,
+                allow_custom_value=True,
+            )
+            
+            with gr.Row():
+                download_btn = gr.Button("📥 Download & List Episodes", variant="secondary", size="lg")
+            
+            download_status = gr.Textbox(
+                label="Download Status",
+                lines=2,
+                interactive=False,
+            )
+        
+        with gr.Tab("Step 2: Select Episodes to Delete"):
+            gr.Markdown("Select the episodes you want to **remove** from the dataset.")
+            
+            episodes_selector = gr.CheckboxGroup(
+                label="Episodes (select to delete)",
+                choices=[],
+                interactive=True,
+            )
+            
+            selected_count = gr.Markdown("*No episodes selected*")
+        
+        with gr.Tab("Step 3: Delete & Upload"):
+            gr.Markdown("Configure the destination and start the deletion process.")
+            
+            dest_repo_input = gr.Textbox(
+                label="Destination Repository ID",
+                placeholder="your-org/cleaned-dataset",
+                info="The HuggingFace repo where the cleaned dataset will be uploaded"
+            )
+            
+            with gr.Row():
+                execute_btn = gr.Button(
+                    "🚀 Delete Episodes & Upload",
+                    variant="primary",
+                    size="lg",
+                )
+            
+            progress_log = gr.Textbox(
+                label="Progress Log",
+                lines=25,
+                interactive=False,
+                max_lines=30,
+            )
+    
+    # Event handlers
+    def update_search_results(query):
+        results = search_datasets_fn(query)
+        return gr.update(choices=results, value=None)
+    
+    search_btn.click(
+        update_search_results,
+        inputs=search_box,
+        outputs=dataset_dropdown,
+    )
+    
+    download_btn.click(
+        download_and_list_episodes,
+        inputs=dataset_dropdown,
+        outputs=[download_status, episodes_selector],
+    )
+    
+    def update_selected_count(selected):
+        if not selected:
+            return "*No episodes selected*"
+        return f"**{len(selected)} episode(s) selected for deletion**"
+    
+    episodes_selector.change(
+        update_selected_count,
+        inputs=episodes_selector,
+        outputs=selected_count,
+    )
+    
+    execute_btn.click(
+        delete_episodes_stream,
+        inputs=[dataset_dropdown, episodes_selector, dest_repo_input],
+        outputs=progress_log,
+    )
+
 
-def greet(name):
-    return "Hello " + name + "!!"
+if __name__ == "__main__":
+    demo.launch()
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()

delete_episodes.py

@@ -0,0 +1,332 @@
+import os
+import re
+import sys
+import glob
+import json
+import logging
+import shutil
+import subprocess
+from pathlib import Path
+from typing import List, Optional, Tuple
+
+from huggingface_hub import snapshot_download, upload_folder, create_repo
+import pandas as pd
+
+
+logger = logging.getLogger(__name__)
+if not logger.handlers:
+    logging.basicConfig(level=logging.INFO, format="[%(levelname)s] %(message)s")
+
+
+def _enable_hf_transfer():
+    """Enable hf_transfer acceleration if the package is installed"""
+    if os.environ.get("HF_HUB_ENABLE_HF_TRANSFER") != "1":
+        os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
+        logger.info("Enabled hf_transfer acceleration (HF_HUB_ENABLE_HF_TRANSFER=1)")
+
+
+def download_dataset(
+    repo_id: str,
+    local_dir: str,
+    hf_token: Optional[str] = None,
+) -> str:
+    """Download a Hugging Face dataset by repo_id.
+    
+    Returns the local directory path.
+    """
+    _enable_hf_transfer()
+    
+    local_path = Path(local_dir)
+    local_path.mkdir(parents=True, exist_ok=True)
+    
+    logger.info(f"Downloading dataset '{repo_id}' to '{local_dir}' ...")
+    
+    path = snapshot_download(
+        repo_id=repo_id,
+        repo_type="dataset",
+        token=hf_token,
+        local_dir=str(local_dir),
+        local_dir_use_symlinks=False,
+    )
+    
+    logger.info(f"Downloaded: {repo_id} -> {path}")
+    return str(local_path)
+
+
+def check_v2_format(dataset_path: str) -> bool:
+    """Check if dataset is in v2.0 format"""
+    info_path = os.path.join(dataset_path, "meta", "info.json")
+    
+    if not os.path.exists(info_path):
+        raise ValueError(f"Error: {info_path} does not exist")
+    
+    with open(info_path, "r") as f:
+        try:
+            info = json.load(f)
+            if "codebase_version" not in info:
+                raise ValueError(f"Error: {info_path} is not a valid v2.0 dataset")
+            elif info["codebase_version"] != "v2.0":
+                raise ValueError(
+                    f"Error: {info_path} is not a v2.0 dataset, found {info['codebase_version']}"
+                )
+            return True
+        except json.JSONDecodeError:
+            raise ValueError(f"Error: {info_path} is not a valid JSON file")
+
+
+def list_episodes(dataset_path: str) -> List[int]:
+    """List all episode numbers in the dataset"""
+    parquets_folder = os.path.join(dataset_path, "data", "chunk-000")
+    
+    if not os.path.exists(parquets_folder):
+        return []
+    
+    parquet_files = glob.glob(os.path.join(parquets_folder, "episode_*.parquet"))
+    
+    episode_numbers = []
+    for file in parquet_files:
+        match = re.search(r"episode_(\d+)\.parquet", file)
+        if match:
+            episode_numbers.append(int(match.group(1)))
+    
+    return sorted(episode_numbers)
+
+
+def delete_ds_store(dataset_path: str):
+    """Delete all .DS_Store files in the given dataset path and its subdirectories"""
+    logger.info("Deleting .DS_Store files...")
+    ds_store_files = glob.glob(
+        os.path.join(dataset_path, "**", ".DS_Store"), recursive=True
+    )
+    
+    if not ds_store_files:
+        logger.info("No .DS_Store files found")
+        return
+    
+    for file in ds_store_files:
+        os.remove(file)
+        logger.info(f"Deleted {file}")
+    
+    logger.info(".DS_Store files deleted")
+
+
+def delete_episode_files(dataset_path: str, indexes: List[int]):
+    """Delete parquet and video files for specified episode indexes"""
+    parquets_folder = os.path.join(dataset_path, "data", "chunk-000")
+    videos_folder = os.path.join(dataset_path, "videos", "chunk-000")
+    
+    # Delete parquet files
+    logger.info("Deleting parquet files...")
+    parquet_files = glob.glob(os.path.join(parquets_folder, "*.parquet"))
+    for index in indexes:
+        for file in parquet_files:
+            if f"episode_{index:06d}.parquet" in file:
+                os.remove(file)
+                logger.info(f"Deleted file {file}")
+    
+    # Delete video files
+    logger.info("Deleting video files...")
+    if os.path.exists(videos_folder):
+        video_folders = os.listdir(videos_folder)
+        for index in indexes:
+            for folder in video_folders:
+                video_files = glob.glob(
+                    os.path.join(videos_folder, folder, f"episode_{index:06d}.mp4")
+                )
+                for video_file in video_files:
+                    os.remove(video_file)
+                    logger.info(f"Deleted file {video_file}")
+
+
+def process_parquet_files(dataset_path: str):
+    """Process all parquet files by correcting the episode_index column"""
+    parquets_folder = os.path.join(dataset_path, "data", "chunk-000")
+    videos_folder = os.path.join(dataset_path, "videos", "chunk-000")
+    
+    logger.info("Processing parquet files...")
+    parquet_files = glob.glob(os.path.join(parquets_folder, "episode_*.parquet"))
+    
+    if not parquet_files:
+        logger.info(f"No parquet files found in {parquets_folder}")
+        return
+    
+    logger.info(f"Found {len(parquet_files)} parquet files to process")
+    
+    # Order files by episode number
+    parquet_files.sort(
+        key=lambda x: int(re.search(r"episode_(\d+)\.parquet", x).group(1))
+    )
+    
+    # Check if episode numbers are continuous
+    episode_numbers = [
+        int(re.search(r"episode_(\d+)\.parquet", file).group(1))
+        for file in parquet_files
+    ]
+    episode_numbers.sort()
+    
+    # Get video folders if they exist
+    video_folders = []
+    if os.path.exists(videos_folder):
+        video_folders = os.listdir(videos_folder)
+    
+    if episode_numbers != list(range(len(episode_numbers))):
+        logger.info(
+            "Episode numbers are not continuous or starting from 0. Renaming files and videos..."
+        )
+        for i, file in enumerate(parquet_files):
+            new_episode_number = i
+            new_file = os.path.join(
+                parquets_folder, f"episode_{new_episode_number:06d}.parquet"
+            )
+            os.rename(file, new_file)
+            logger.info(f"Renamed {file} to {new_file}")
+            
+            # Rename corresponding video files
+            for folder in video_folders:
+                video_file = os.path.join(
+                    videos_folder, folder, f"episode_{episode_numbers[i]:06d}.mp4"
+                )
+                new_video_file = os.path.join(
+                    videos_folder, folder, f"episode_{new_episode_number:06d}.mp4"
+                )
+                if os.path.exists(video_file):
+                    os.rename(video_file, new_video_file)
+                    logger.info(f"Renamed {video_file} to {new_video_file}")
+        
+        # Update list after renaming
+        parquet_files = glob.glob(os.path.join(parquets_folder, "episode_*.parquet"))
+        parquet_files.sort(
+            key=lambda x: int(re.search(r"episode_(\d+)\.parquet", x).group(1))
+        )
+        logger.info("Updated parquet files list after renaming")
+    
+    # Process each parquet file
+    total_index = 0
+    for file_path in parquet_files:
+        filename = os.path.basename(file_path)
+        match = re.search(r"episode_(\d+)\.parquet", filename)
+        
+        if match:
+            episode_number = int(match.group(1))
+            logger.info(f"Processing {filename} - Episode {episode_number}")
+            
+            try:
+                df = pd.read_parquet(file_path, engine="pyarrow")
+                
+                df["episode_index"] = episode_number
+                df["frame_index"] = range(len(df))
+                df["index"] = range(total_index, total_index + len(df))
+                total_index += len(df)
+                
+                df.to_parquet(file_path, index=False)
+                logger.info(f"Successfully updated {filename}")
+                
+            except Exception as e:
+                raise RuntimeError(f"Error processing {filename}: {str(e)}")
+        else:
+            logger.info(f"Skipping {filename} - doesn't match expected pattern")
+    
+    logger.info("Parquet processing complete")
+
+
+def run_stats_computation(dataset_path: str):
+    """Run the lerobot stats computation script"""
+    script_path = "lerobot_stats_compute.py"
+    
+    if not os.path.exists(script_path):
+        logger.warning(f"Stats script '{script_path}' not found, skipping stats computation")
+        return
+    
+    logger.info("Running lerobot_stats_compute.py...")
+    
+    try:
+        subprocess.run(
+            ["uv", "run", script_path, "--dataset-path", dataset_path],
+            check=True,
+        )
+        logger.info(f"Successfully executed {script_path}")
+    except subprocess.CalledProcessError as e:
+        logger.warning(f"Error executing stats script: {str(e)}")
+    except FileNotFoundError:
+        logger.warning("uv not found, skipping stats computation")
+
+
+def delete_episodes_and_repair(
+    dataset_path: str,
+    episode_indexes: List[int],
+    run_stats: bool = True,
+) -> str:
+    """Delete specified episodes and repair the dataset.
+    
+    Args:
+        dataset_path: Path to the dataset
+        episode_indexes: List of episode indexes to delete
+        run_stats: Whether to run stats computation after repair
+        
+    Returns:
+        Path to the repaired dataset
+    """
+    if not episode_indexes:
+        raise ValueError("No episode indexes provided for deletion")
+    
+    # Check v2.0 format
+    check_v2_format(dataset_path)
+    
+    logger.info(f"Deleting episodes: {episode_indexes}")
+    
+    # Delete .DS_Store files
+    delete_ds_store(dataset_path)
+    
+    # Delete episode files
+    delete_episode_files(dataset_path, episode_indexes)
+    
+    # Process and repair remaining parquet files
+    process_parquet_files(dataset_path)
+    
+    # Run stats computation
+    if run_stats:
+        run_stats_computation(dataset_path)
+    
+    logger.info("Episode deletion and repair complete")
+    return dataset_path
+
+
+def upload_dataset(
+    local_dir: str,
+    dest_repo_id: str,
+    hf_token: Optional[str] = None,
+    commit_message: Optional[str] = None,
+    private: bool = False,
+) -> str:
+    """Upload a local dataset folder to a destination HF dataset repo.
+    
+    Returns the repo URL/identifier.
+    """
+    if not dest_repo_id:
+        raise ValueError("dest_repo_id must be provided")
+    
+    token = hf_token or os.environ.get("HF_TOKEN")
+    create_repo(
+        repo_id=dest_repo_id,
+        repo_type="dataset",
+        private=private,
+        exist_ok=True,
+        token=token,
+    )
+    
+    _enable_hf_transfer()
+    msg = commit_message or "Updated dataset after episode deletion"
+    logger.info(f"Uploading '{local_dir}' to '{dest_repo_id}' (private={private}) ...")
+    
+    upload_folder(
+        repo_id=dest_repo_id,
+        repo_type="dataset",
+        folder_path=local_dir,
+        path_in_repo=".",
+        commit_message=msg,
+        token=token,
+    )
+    
+    logger.info(f"Uploaded to: {dest_repo_id}")
+    return dest_repo_id
+

requirements.txt

@@ -0,0 +1,73 @@
+aiofiles==24.1.0
+aiohappyeyeballs==2.6.1
+aiohttp==3.12.15
+aiosignal==1.4.0
+annotated-types==0.7.0
+anyio==4.11.0
+attrs==25.3.0
+Brotli==1.1.0
+certifi==2025.8.3
+charset-normalizer==3.4.3
+click==8.3.0
+datasets==4.1.1
+dill==0.4.0
+fastapi==0.117.1
+ffmpy==0.6.1
+filelock==3.19.1
+frozenlist==1.7.0
+fsspec==2025.9.0
+gradio==5.47.0
+gradio_client==1.13.2
+groovy==0.1.2
+h11==0.16.0
+hf-xet==1.1.10
+hf_transfer==0.1.9
+httpcore==1.0.9
+httpx==0.28.1
+huggingface-hub==0.35.1
+idna==3.10
+inquirerpy==0.3.4
+Jinja2==3.1.6
+markdown-it-py==4.0.0
+MarkupSafe==3.0.2
+mdurl==0.1.2
+multidict==6.6.4
+multiprocess==0.70.16
+numpy==2.2.6
+orjson==3.11.3
+packaging==25.0
+pandas==2.3.2
+pfzy==0.3.4
+pillow==11.3.0
+prompt_toolkit==3.0.52
+propcache==0.3.2
+pyarrow==21.0.0
+pydantic==2.11.9
+pydantic_core==2.33.2
+pydub==0.25.1
+Pygments==2.19.2
+python-dateutil==2.9.0.post0
+python-multipart==0.0.20
+pytz==2025.2
+PyYAML==6.0.2
+requests==2.32.5
+rich==14.1.0
+ruff==0.13.1
+safehttpx==0.1.6
+semantic-version==2.10.0
+shellingham==1.5.4
+six==1.17.0
+sniffio==1.3.1
+starlette==0.48.0
+tomlkit==0.13.3
+tqdm==4.67.1
+typer==0.19.2
+typing-inspection==0.4.1
+typing_extensions==4.15.0
+tzdata==2025.2
+urllib3==2.5.0
+uvicorn==0.37.0
+wcwidth==0.2.14
+websockets==15.0.1
+xxhash==3.5.0
+yarl==1.20.1