add decode logging

scripts/convert_videos_to_images.py

@@ -0,0 +1,74 @@
+#!/usr/bin/env python
+
+"""
+Convert video dataset to image dataset for faster training.
+This pre-extracts all frames from MP4 files to PNG images.
+"""
+
+import argparse
+from pathlib import Path
+import logging
+import shutil
+
+def convert_dataset_videos_to_images(repo_id: str, root: str | None = None):
+    """Convert all videos in a LeRobot dataset to individual image files."""
+    from lerobot.datasets.lerobot_dataset import LeRobotDataset
+    from lerobot.datasets.video_utils import decode_video_frames
+    import torch
+    
+    # Load dataset
+    dataset = LeRobotDataset(repo_id, root=root, download_videos=True)
+    
+    total_frames_processed = 0
+    
+    for ep_idx in range(dataset.meta.total_episodes):
+        logging.info(f"Processing episode {ep_idx}/{dataset.meta.total_episodes}")
+        
+        for vid_key in dataset.meta.video_keys:
+            video_path = dataset.root / dataset.meta.get_video_file_path(ep_idx, vid_key)
+            
+            if not video_path.exists():
+                logging.warning(f"Video not found: {video_path}")
+                continue
+                
+            # Create image directory  
+            img_dir = dataset.root / f"images/chunk-{dataset.meta.get_episode_chunk(ep_idx)}/{vid_key}"
+            img_dir.mkdir(parents=True, exist_ok=True)
+            
+            # Decode all frames from video
+            # Get episode length to decode all frames
+            ep_length = dataset.meta.episodes[ep_idx]["length"]
+            timestamps = [i / dataset.fps for i in range(ep_length)]
+            
+            try:
+                frames = decode_video_frames(video_path, timestamps, dataset.tolerance_s, dataset.video_backend)
+                
+                # Save each frame as PNG
+                for i, frame in enumerate(frames.squeeze(0)):
+                    img_path = img_dir / f"episode_{ep_idx:06d}_{i:06d}.png"
+                    # Convert tensor to PIL and save
+                    import torchvision.transforms as T
+                    to_pil = T.ToPILImage()
+                    pil_frame = to_pil(frame)
+                    pil_frame.save(img_path)
+                    
+                total_frames_processed += len(frames.squeeze(0))
+                logging.info(f"  Extracted {len(frames.squeeze(0))} frames to {img_dir}")
+                
+            except Exception as e:
+                logging.error(f"Failed to process {video_path}: {e}")
+                continue
+    
+    logging.info(f"Conversion complete! Processed {total_frames_processed} total frames")
+    logging.info(f"You can now use download_videos=False to use the extracted images")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Convert LeRobot video dataset to images")
+    parser.add_argument("repo_id", help="Dataset repo ID (e.g., 'kenmacken/record-test-2')")
+    parser.add_argument("--root", help="Local root directory", default=None)
+    
+    args = parser.parse_args()
+    
+    logging.basicConfig(level=logging.INFO)
+    
+    convert_dataset_videos_to_images(args.repo_id, args.root)