時間がかかったけど、なんとかｚarr化できたのでメモ。
ドキュメントはここ
https://github.com/real-stanford/diffusion_policy?tab=readme-ov-file#replaybuffer
data/pusht_cchi_v7_replay.zarr
├── data
│ ├── action (25650, 2) float32
│ ├── img (25650, 96, 96, 3) float32
│ ├── keypoint (25650, 9, 2) float32
│ ├── n_contacts (25650, 1) float32
│ └── state (25650, 5) float32
└── meta
└── episode_ends (206,) int64
この構成のzarrデータを作る必要がある。
dataの下に入力データを並べ、metaの下にepisode_endsの配列を用意する。
入力データは全部まとめて一連の配列にする。配列の区切り（エピソードの区切り）をepisode_endsに入れておく。
zarr.saveだけではこの構造で保存できず。
store = zarr.DirectoryStore(OUTPUT_FILE)
root = zarr.group(store=store, overwrite=True)
こうやってから、個別にroot.create_datasetで上手くデータ作れました。
import os
import glob
import zarr
from PIL import Image
import numpy as np
ROOTDIR = 'dofbot_data/OK'
OUTPUT_FILE = 'dofbot.zarr'
def create_state(root_dir:str) -> zarr.core.Array:
pose_list = []
for d in sorted(os.listdir(root_dir)):
pose_csv = os.path.join(root_dir, d, 'pose.csv')
with open(pose_csv, 'r') as f:
for line in f:
pose_words = line.split(',')
pose = [float(w) for w in pose_words]
pose_list.append(pose)
z = zarr.array(pose_list, chunks=(1, len(pose_list[0])), dtype='float32')
return z
def create_images(root_dir:str) -> zarr.core.Array:
images = []
for d in sorted(os.listdir(root_dir)):
files = sorted(glob.glob(os.path.join(root_dir, d, '*.png')))
for f in files:
img = Image.open(f)
img = img.resize((480, 480))
img_arr = np.array(img).astype('uint8')
images.append(img_arr)
z = zarr.array(images, chunks=(1, 480, 480, 3), dtype='float32')
return z
def calc_episode_ends(dir: str) -> list:
dirs = sorted(os.listdir(dir))
episode_ends = []
ei = 0
for d in dirs:
files = glob.glob(os.path.join(dir, d, '*.png'))
ei += len(files)
episode_ends.append(ei)
return episode_ends
def create_action(states: zarr.core.Array, episode_ends:list) -> zarr.core.Array:
actions = []
for idx in range(len(states)):
if idx + 1 in episode_ends:
action = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
else:
action_now = states[idx]
action_next = states[idx+1]
action = [action_next[0] - action_now[0],
action_next[1] - action_now[1],
action_next[2] - action_now[2],
action_next[3] - action_now[3],
action_next[4] - action_now[4],
action_next[5] - action_now[5]]
actions.append(action)
z = zarr.array(actions, chunks=(1,), dtype='float32')
return z
def main():
episode_ends = calc_episode_ends(ROOTDIR)
states = create_state(ROOTDIR)
actions = create_action(states, episode_ends)
images = create_images(ROOTDIR)
store = zarr.DirectoryStore(OUTPUT_FILE)
root = zarr.group(store=store, overwrite=True)
root.create_dataset('meta/episode_ends', data=episode_ends)
root.create_dataset('data/state', data=states)
root.create_dataset('data/action', data=actions)
root.create_dataset('data/img', data=images)
print(root.info)
print('saved', OUTPUT_FILE)
if __name__ == '__main__':
main()
