Diffusion Policyの学習データを覗いてみる

Diffusion Policyを動かそうとしています。

github.com

学習データがhttps://diffusion-policy.cs.columbia.edu/data/training/からダウンロードできます。 zarrというフォーマットから画像を取り出して確認してみました。

ソース

import zarr  
from PIL import Image  
  
#https://diffusion-policy.cs.columbia.edu/data/training/  
z = zarr.open('data/pusht_cchi_v7_replay.zarr')  
  
zimg = z['data/img']  
  
first_img = zimg[0].astype('uint8')  
pil_img = Image.fromarray(first_img)  
pil_img.save('first.png')  
  
print(zimg.info)

ダウンロードしたデータの位置を合わせて実行するとこう表示されます。

Name               : /data/img
Type               : zarr.core.Array
Data type          : float32
Shape              : (25650, 96, 96, 3)
Chunk shape        : (161, 96, 96, 3)
Order              : C
Read-only          : False
Compressor         : Blosc(cname='zstd', clevel=5, shuffle=BITSHUFFLE,
                   : blocksize=0)
Store type         : zarr.storage.DirectoryStore
No. bytes          : 2836684800 (2.6G)
No. bytes stored   : 29791685 (28.4M)
Storage ratio      : 95.2
Chunks initialized : 160/160

ポイントとしては、zimg = z['data/img'] とディレクトリのようにして画像ファイルの塊へアクセスすることと、uint8に変換することくらいです。

最初の画像はこうなっていて、puth Tのデータになってます。