Lerobot-Dataset-Ruby 0.1.0

data/lib/Lerobot/Dataset/Ruby/writer.rb

@@ -0,0 +1,362 @@
+# frozen_string_literal: true
+
+require 'json'
+require 'fileutils'
+require 'parquet'
+
+module Lerobot
+  module Dataset
+    module Ruby
+      class Writer
+        attr_reader :output_dir, :fps
+
+        def initialize(output_dir:, fps: 30, tasks: nil)
+          @output_dir = File.expand_path(output_dir)
+          @fps = fps
+          @episodes = []
+          @all_frames = []
+          @task_map = {}
+          @episode_counter = 0
+          @next_video_frame_index = 0
+          @recording = false
+
+          Array(tasks).each { |task| task_index(task) }
+
+          # v3ディレクトリ構造
+          FileUtils.mkdir_p(File.join(@output_dir, 'meta'))
+          FileUtils.mkdir_p(File.join(@output_dir, 'data', 'chunk-000'))
+          FileUtils.mkdir_p(File.join(@output_dir, 'videos', 'rgb', 'chunk-000'))
+          FileUtils.mkdir_p(File.join(@output_dir, 'meta', 'episodes', 'chunk-000'))
+        end
+
+        # 録画開始（内部でLibopencvを呼び出す）
+        def start_recording(path = nil, fps = nil)
+          begin
+            require 'libopencv'
+          rescue LoadError
+            require_relative '../../../../ext/opencv/libopencv'
+          end
+          path ||= File.join(@output_dir, 'videos', 'rgb', 'chunk-000', 'file-000.mp4')
+          Libopencv.start_recording(path, fps || @fps)
+          image_bytes = Libopencv.take_picture
+          Libopencv.write_frame(image_bytes)
+          @recording = true
+          puts "録画開始: #{path}"
+        end
+
+        # 録画停止（内部でLibopencvを呼び出す）
+        def stop_recording
+          if @recording
+            Libopencv.stop_recording
+            @recording = false
+            puts "録画を停止しました"
+          end
+        end
+
+        # 新しいエピソードを開始
+        def start_episode(task: "Default task")
+          @current_episode = {
+            episode_index: @episode_counter,
+            task: task,
+            from_index: @all_frames.length
+          }
+          @episode_counter += 1
+          task_index(task)
+        end
+
+        # フレームを追加（現在のエピソードに）
+        def add_frame(action:, state:, image_bytes: nil, timestamp: nil)
+          unless @current_episode
+            raise RuntimeError, "start_episodeを最初に呼び出してください"
+          end
+
+          episode_frame_index = @all_frames.length - @current_episode[:from_index]
+          frame_data = {
+            index: @all_frames.length,
+            episode_index: @current_episode[:episode_index],
+            frame_index: @next_video_frame_index,
+            task_index: task_index(@current_episode[:task]),
+            action: action,
+            state: state,
+            timestamp: timestamp || episode_frame_index.to_f / @fps
+          }
+
+          if image_bytes
+            # 録画中ならLibopencvに書き込み
+            raise RuntimeError, "録画が開始されていません" unless @recording
+            Libopencv.write_frame(image_bytes)
+          end
+
+          @all_frames << frame_data
+          @next_video_frame_index += 1
+        end
+
+        # エピソードを保存（save: false なら動画だけ残してデータセットから除外）
+        def save_episode(save: true)
+          unless @current_episode
+            return
+          end
+
+          episode = @current_episode
+
+          unless save
+            @all_frames.slice!(episode[:from_index], @all_frames.length - episode[:from_index])
+            @all_frames.each_with_index { |frame, index| frame[:index] = index }
+            @episode_counter -= 1
+            @current_episode = nil
+            return
+          end
+
+          episode_length = @all_frames.length - episode[:from_index]
+          episode_frames = @all_frames[episode[:from_index], episode_length] || []
+          first_video_frame = episode_frames.first ? episode_frames.first[:frame_index] : @next_video_frame_index
+          last_video_frame = episode_frames.last ? episode_frames.last[:frame_index] : first_video_frame - 1
+
+          @episodes << {
+            episode_index: episode[:episode_index],
+            tasks: [episode[:task]],
+            length: episode_length,
+            dataset_from_index: episode[:from_index],
+            dataset_to_index: @all_frames.length,
+            "data/chunk_index" => 0,
+            "data/file_index" => 0,
+            "videos/rgb/chunk_index" => 0,
+            "videos/rgb/file_index" => 0,
+            "videos/rgb/from_timestamp" => first_video_frame.to_f / @fps,
+            "videos/rgb/to_timestamp" => (last_video_frame + 1).to_f / @fps,
+            "meta/episodes/chunk_index" => 0,
+            "meta/episodes/file_index" => 0,
+            frame_start: episode[:from_index],
+            frame_end: @all_frames.length,
+            task_index: task_index(episode[:task])
+          }
+
+          @current_episode = nil
+        end
+
+        # 最終保存（Parquet書き込み）
+        def save
+          if @current_episode
+            save_episode
+          end
+
+          ensure_default_video_file
+          write_info_json
+          write_tasks_parquet
+          write_stats_json
+          write_episodes_parquet
+          write_data_parquet
+
+          puts "Dataset v3.0 successfully created at: #{@output_dir}"
+        end
+
+        private
+
+        def task_index(task)
+          @task_map[task] ||= @task_map.size
+        end
+
+        def ensure_default_video_file
+          video_dir = File.join(@output_dir, 'videos', 'rgb', 'chunk-000')
+          video_path = File.join(video_dir, 'file-000.mp4')
+          FileUtils.mkdir_p(video_dir)
+
+          sample_candidates = [
+            File.expand_path('../../../../test/test.mp4', __dir__),
+            File.expand_path('../../../../public/test.mp4', __dir__),
+            File.join(__dir__, 'test', 'test.mp4'),
+            File.join(__dir__, 'public', 'test.mp4')
+          ]
+          sample_path = sample_candidates.find { |path| File.exist?(path) }
+
+          if !File.exist?(video_path) && sample_path
+            FileUtils.cp(sample_path, video_path)
+          end
+          video_path
+        end
+
+        def write_info_json
+          action_size = @all_frames.first&.dig(:action)&.size || 2
+          state_size = @all_frames.first&.dig(:state)&.size || 2
+
+          info = {
+            codebase_version: "v3.0",
+            fps: @fps,
+            total_episodes: @episodes.size,
+            total_frames: @all_frames.size,
+            total_tasks: @task_map.size,
+            chunks_size: 1000,
+            data_files_size_in_mb: 100,
+            video_files_size_in_mb: 200,
+            data_path: "data/chunk-{chunk_index:03d}/file-{file_index:03d}.parquet",
+            video_path: "videos/{video_key}/chunk-{chunk_index:03d}/file-{file_index:03d}.mp4",
+            features: {
+              "index" => { "dtype" => "int64", "shape" => [1] },
+              "episode_index" => { "dtype" => "int64", "shape" => [1] },
+              "frame_index" => { "dtype" => "int64", "shape" => [1] },
+              "task_index" => { "dtype" => "int64", "shape" => [1] },
+              "action" => { "dtype" => "float32", "shape" => [action_size] },
+              "state" => { "dtype" => "float32", "shape" => [state_size] },
+              "timestamp" => { "dtype" => "float32", "shape" => [1] },
+              "rgb" => { "dtype" => "video", "shape" => [480, 640, 3], "names" => ["height", "width", "channel"] }
+            },
+            tasks: @task_map.sort_by { |_, index| index }.map(&:first)
+          }
+
+          File.write(File.join(@output_dir, 'meta', 'info.json'), JSON.pretty_generate(info))
+        end
+
+        def write_tasks_parquet
+          parquet_path = File.join(@output_dir, 'meta', 'tasks.parquet')
+          schema = Parquet::Schema.define do
+            field :task_index, :int64, nullable: false
+            field :task, :string, nullable: false
+          end
+
+          rows = @task_map.sort_by { |_, index| index }.map do |task_str, index|
+            [index, task_str]
+          end
+
+          Parquet.write_rows(rows, schema: schema, write_to: parquet_path)
+        end
+
+        def write_stats_json
+          action_size = @all_frames.first&.dig(:action)&.size || 2
+          state_size = @all_frames.first&.dig(:state)&.size || 2
+          total_frames = @all_frames.size
+
+          stats = {
+            "index" => {
+              "mean" => [0],
+              "std" => [1],
+              "min" => [0],
+              "max" => [1],
+              "count" => total_frames
+            },
+            "episode_index" => {
+              "mean" => [0],
+              "std" => [1],
+              "min" => [0],
+              "max" => [1],
+              "count" => total_frames
+            },
+            "frame_index" => {
+              "mean" => [0],
+              "std" => [1],
+              "min" => [0],
+              "max" => [1],
+              "count" => total_frames
+            },
+            "task_index" => {
+              "mean" => [0],
+              "std" => [1],
+              "min" => [0],
+              "max" => [1],
+              "count" => total_frames
+            },
+            "action" => {
+              "mean" => Array.new(action_size, 0.0),
+              "std" => Array.new(action_size, 1.0),
+              "min" => Array.new(action_size, -1.0),
+              "max" => Array.new(action_size, 1.0),
+              "count" => total_frames
+            },
+            "state" => {
+              "mean" => Array.new(state_size, 0.0),
+              "std" => Array.new(state_size, 1.0),
+              "min" => Array.new(state_size, -1.0),
+              "max" => Array.new(state_size, 1.0),
+              "count" => total_frames
+            },
+            "timestamp" => {
+              "mean" => [0.0],
+              "std" => [1.0],
+              "min" => [0.0],
+              "max" => [1.0],
+              "count" => total_frames
+            }
+          }
+
+          File.write(File.join(@output_dir, 'meta', 'stats.json'), JSON.pretty_generate(stats))
+        end
+
+        def write_episodes_parquet
+          parquet_path = File.join(@output_dir, 'meta', 'episodes', 'chunk-000', 'file-000.parquet')
+
+          schema = Parquet::Schema.define do
+            field :episode_index, :int32, nullable: false
+            field :tasks, :list, item: :string
+            field :length, :int32, nullable: false
+            field :dataset_from_index, :int32, nullable: false
+            field :dataset_to_index, :int32, nullable: false
+            field :"data/chunk_index", :int32, nullable: false
+            field :"data/file_index", :int32, nullable: false
+            field :"videos/rgb/chunk_index", :int32, nullable: false
+            field :"videos/rgb/file_index", :int32, nullable: false
+            field :"videos/rgb/from_timestamp", :float
+            field :"videos/rgb/to_timestamp", :float
+            field :"meta/episodes/chunk_index", :int32, nullable: false
+            field :"meta/episodes/file_index", :int32, nullable: false
+            field :frame_start, :int32, nullable: false
+            field :frame_end, :int32, nullable: false
+            field :task_index, :int64, nullable: false
+          end
+
+          rows = @episodes.map do |ep|
+            [
+              ep[:episode_index],
+              ep[:tasks],
+              ep[:length],
+              ep[:dataset_from_index],
+              ep[:dataset_to_index],
+              ep["data/chunk_index"],
+              ep["data/file_index"],
+              ep["videos/rgb/chunk_index"],
+              ep["videos/rgb/file_index"],
+              ep["videos/rgb/from_timestamp"],
+              ep["videos/rgb/to_timestamp"],
+              ep["meta/episodes/chunk_index"],
+              ep["meta/episodes/file_index"],
+              ep[:frame_start],
+              ep[:frame_end],
+              ep[:task_index]
+            ]
+          end
+
+          Parquet.write_rows(rows, schema: schema, write_to: parquet_path)
+        end
+
+        def write_data_parquet
+          parquet_path = File.join(@output_dir, 'data', 'chunk-000', 'file-000.parquet')
+
+          schema = Parquet::Schema.define do
+            field :index, :int64, nullable: false
+            field :episode_index, :int64, nullable: false
+            field :frame_index, :int64, nullable: false
+            field :task_index, :int64, nullable: false
+            field :action, :list, item: :float
+            field :state, :list, item: :float
+            field :timestamp, :float
+          end
+
+          rows = @all_frames.map do |frame|
+            [
+              frame[:index],
+              frame[:episode_index],
+              frame[:frame_index],
+              frame[:task_index],
+              frame[:action],
+              frame[:state],
+              frame[:timestamp]
+            ]
+          end
+
+          Parquet.write_rows(rows, schema: schema, write_to: parquet_path)
+        end
+      end
+    end
+  end
+end
+
+# 後方互換性のためのエイリアス
+LeRobotDatasetWriter = Lerobot::Dataset::Ruby::Writer

data/lib/Lerobot/Dataset/Ruby.rb

@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+require_relative "Ruby/version"
+require_relative "Ruby/reader"
+require_relative "Ruby/writer"
+
+module Lerobot
+  module Dataset
+    module Ruby
+      class Error < StandardError; end
+    end
+  end
+end

data/sig/Lerobot/Dataset/Ruby.rbs

@@ -0,0 +1,8 @@
+module Lerobot
+  module Dataset
+    module Ruby
+      VERSION: String
+      # See the writing guide of rbs: https://github.com/ruby/rbs#guides
+    end
+  end
+end

data/test/minimal_lerobot_check.rb

@@ -0,0 +1,49 @@
+#!/usr/bin/env ruby
+# frozen_string_literal: true
+
+require 'fileutils'
+require_relative '../lerobot_dataset_writer'
+require_relative '../ext/opencv/libopencv'
+
+root = File.expand_path('../output_dataset/robot_dataset', __dir__)
+FileUtils.rm_rf(root)
+
+writer = LeRobotDatasetWriter.new(output_dir: root, fps: 30)
+writer.start_episode(task: 'Grab the black cube')
+
+3.times do |i|
+  writer.add_frame(
+    action: [0.1 * (i + 1), 0.2 * (i + 1)],
+    state: [1.0, 0.5],
+    timestamp: i.to_f / 30.0
+  )
+end
+
+writer.save_episode
+writer.save
+puts "Ruby dataset generated at: #{root}"
+
+python_code = <<~PY
+  import json
+  from pathlib import Path
+  from lerobot.datasets.lerobot_dataset import LeRobotDataset
+
+  root = Path(r"#{root}")
+  print("root_exists", root.exists())
+  ds = LeRobotDataset(repo_id="lerobot_min_check", root=root)
+  print("len", len(ds))
+  print("tasks", ds.meta.tasks.to_dict())
+  sample = ds[0]
+  print("sample_keys", sorted(sample.keys()))
+  print("timestamp", sample["timestamp"])
+  print("action", sample["action"])
+  print("state", sample["state"])
+PY
+
+venv_activate = File.expand_path('../.venv/bin/activate', __dir__)
+system(
+  'bash', '-lc',
+  "cd #{File.expand_path('..', __dir__)} && . \"#{venv_activate}\" && python - <<'PY'\n#{python_code}\nPY"
+)
+
+exit($?.exitstatus)

data/test/sample.rb

@@ -0,0 +1,56 @@
+require_relative '../lerobot_dataset_writer'
+require_relative '../ext/opencv/libopencv'
+
+# データセット書き込みインスタンス作成
+writer = LeRobotDatasetWriter.new(output_dir: File.expand_path('../output_dataset/robot_dataset', __dir__), fps: 30)
+
+# カメラをオープン（外部からソースを指定）
+puts "カメラをオープン中..."
+# Libopencv.open(0)
+# Libopencv.open_mp4('./public/test.mp4')
+Libopencv.open_mp4(File.expand_path('../test/test.mp4', __dir__))
+
+# 録画開始（ファイル名はWriterのデフォルトを使用）
+writer.start_recording
+
+# 10エピソード
+10.times do |j|
+
+  # エピソード開始（タスク指定）
+  writer.start_episode(task: "Grab the black cube")
+
+  # 10フレーム撮影・録画・データ保存
+  10.times do |i|
+    # カメラから画像取得
+    image_bytes = Libopencv.take_picture
+    
+    if image_bytes
+      # LeRobot データセットにも保存（v3形式）
+      writer.add_frame(
+        action: [0.1 * i, 0.2 * i],
+        state: [1.0, 0.5],
+        image_bytes: image_bytes,
+        timestamp: i.to_f / 30
+      )
+      
+      puts "フレーム #{i} を保存しました"
+    else
+      puts "フレーム #{i} の取得に失敗"
+    end
+  end
+
+  # エピソード保存
+  writer.save_episode
+
+end
+
+# 録画停止（Libopencvの呼び出しはWriter内部で実行）
+writer.stop_recording
+
+# ウィンドウ破棄
+Libopencv.destroy_window
+
+# データセット保存（Parquet書き込み）
+writer.save
+
+puts "完了！"

data/test/test_dataset.py

@@ -0,0 +1,189 @@
+#!/usr/bin/env python3
+"""Validate a local LeRobotDataset v3.0 directory.
+
+Usage:
+    python3 test_dataset.py output_dataset/my_first_robot_dataset
+
+If the ``lerobot`` package is installed, the script also tries to load the
+ dataset through LeRobotDataset. PyArrow and OpenCV are used for independent
+ structural checks.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import pyarrow.parquet as pq
+
+try:
+    import cv2
+except ImportError:
+    cv2 = None
+
+
+class DatasetValidationError(Exception):
+    """Raised when the dataset does not satisfy the expected v3 invariants."""
+
+
+def load_json(path: Path) -> dict:
+    if not path.is_file():
+        raise DatasetValidationError(f"Missing file: {path}")
+    with path.open(encoding="utf-8") as stream:
+        return json.load(stream)
+
+
+def validate_tasks(dataset_root: Path) -> int:
+    tasks_parquet_path = dataset_root / "meta" / "tasks.parquet"
+    if not tasks_parquet_path.is_file():
+        raise DatasetValidationError(f"Missing file: {tasks_parquet_path}")
+
+    table = pq.read_table(tasks_parquet_path)
+    if "task_index" not in table.column_names and (
+        table.schema.pandas_metadata is None or "task_index" not in str(table.schema.pandas_metadata)
+    ):
+        raise DatasetValidationError("meta/tasks.parquet missing task_index column")
+    if len(table) == 0:
+        raise DatasetValidationError("meta/tasks.parquet contains no tasks")
+    return len(table)
+
+
+def validate_dataset(dataset_root: Path) -> dict:
+    info = load_json(dataset_root / "meta" / "info.json")
+    if info.get("codebase_version") != "v3.0":
+        raise DatasetValidationError(
+            f"Expected codebase_version v3.0, got {info.get('codebase_version')!r}"
+        )
+
+    required_paths = (
+        dataset_root / "data" / "chunk-000" / "file-000.parquet",
+        dataset_root / "meta" / "episodes" / "chunk-000" / "file-000.parquet",
+    )
+    for path in required_paths:
+        if not path.is_file():
+            raise DatasetValidationError(f"Missing file: {path}")
+
+    task_count = validate_tasks(dataset_root)
+    data_table = pq.read_table(required_paths[0])
+    episode_table = pq.read_table(required_paths[1])
+    data_columns = set(data_table.column_names)
+    episode_columns = set(episode_table.column_names)
+
+    required_data_columns = {"episode_index", "frame_index", "action", "state", "timestamp"}
+    required_episode_columns = {"episode_index", "tasks", "length"}
+    missing_data = required_data_columns - data_columns
+    missing_episodes = required_episode_columns - episode_columns
+    if missing_data:
+        raise DatasetValidationError(f"Missing data columns: {sorted(missing_data)}")
+    if missing_episodes:
+        raise DatasetValidationError(
+            f"Missing episode columns: {sorted(missing_episodes)}"
+        )
+
+    frame_count = data_table.num_rows
+    episode_count = episode_table.num_rows
+    declared_frames = info.get("total_frames")
+    declared_episodes = info.get("total_episodes")
+    if declared_frames != frame_count:
+        raise DatasetValidationError(
+            f"total_frames={declared_frames} but data contains {frame_count} rows"
+        )
+    if declared_episodes != episode_count:
+        raise DatasetValidationError(
+            f"total_episodes={declared_episodes} but episode metadata contains "
+            f"{episode_count} rows"
+        )
+
+    lengths = episode_table.column("length").to_pylist()
+    if sum(lengths) != frame_count:
+        raise DatasetValidationError(
+            f"Episode lengths sum to {sum(lengths)}, but data contains {frame_count} rows"
+        )
+
+    video_relative_path = (
+        info.get("data_files", {}).get("videos", {}).get("rgb")
+        or (
+            info.get("video_path")
+            and info["video_path"].format(
+                video_key="rgb", chunk_index=0, file_index=0
+            )
+        )
+        or "videos/rgb/chunk-000/file-000.mp4"
+    )
+    if not video_relative_path:
+        raise DatasetValidationError("info.json does not define the rgb video path")
+    video_path = dataset_root / video_relative_path
+    if not video_path.is_file() or video_path.stat().st_size == 0:
+        raise DatasetValidationError(f"Missing or empty video: {video_path}")
+
+    video_frames = None
+    video_fps = None
+    if cv2 is not None:
+        capture = cv2.VideoCapture(str(video_path))
+        if not capture.isOpened():
+            raise DatasetValidationError(f"OpenCV cannot open video: {video_path}")
+        video_frames = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
+        video_fps = capture.get(cv2.CAP_PROP_FPS)
+        capture.release()
+        if video_frames < frame_count:
+            raise DatasetValidationError(
+                f"Video contains {video_frames} frames, but data contains {frame_count} rows"
+            )
+        expected_fps = float(info.get("fps", 0))
+        if expected_fps > 0 and abs(video_fps - expected_fps) > 0.5:
+            raise DatasetValidationError(
+                f"Video FPS is {video_fps}, expected approximately {expected_fps}"
+            )
+
+    return {
+        "episodes": episode_count,
+        "frames": frame_count,
+        "tasks": task_count,
+        "video_frames": video_frames,
+        "video_fps": video_fps,
+    }
+
+
+def validate_with_lerobot(dataset_root: Path) -> str:
+    try:
+        from lerobot.datasets.lerobot_dataset import LeRobotDataset
+    except ImportError:
+        return "LeRobotDataset import skipped: lerobot is not installed"
+
+    try:
+        dataset = LeRobotDataset(repo_id=dataset_root.name, root=str(dataset_root))
+        if len(dataset) == 0:
+            raise DatasetValidationError("LeRobotDataset loaded zero samples")
+        return f"LeRobotDataset loaded successfully ({len(dataset)} samples)"
+    except Exception as error:
+        raise DatasetValidationError(f"LeRobotDataset load failed: {error}") from error
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "dataset_root",
+        nargs="?",
+        default="output_dataset/robot_dataset",
+        type=Path,
+    )
+    args = parser.parse_args()
+    dataset_root = args.dataset_root.resolve()
+
+    try:
+        result = validate_dataset(dataset_root)
+        lerobot_result = validate_with_lerobot(dataset_root)
+    except DatasetValidationError as error:
+        print(f"INVALID: {error}", file=sys.stderr)
+        return 1
+
+    print(f"VALID: {dataset_root}")
+    print(json.dumps(result, indent=2))
+    print(lerobot_result)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())