PyTorch DataLoader で tarを展開せずに読込む

tarファイル

Deep Learning で用いる公開データセットには tar 形式(tar.gzなど含む)のものが存在する.
データ分析やデバッグなどを行うという観点では1枚1枚の画像が見えた方がよいとは思う.
だが tar にしているということはデータ量が多いということであり、
それはつまり、ストレージに優しくないことを指す.
そのため tar のまま扱えた方が取り回しがよいのではないかと考えた.

画像ファイルが格納された tar があるとして、それを展開せずに利用する例を書いた.

import cv2
import numpy as np
import tarfile

arc = tarfile.open('test.tar')
print(arc.getnames())
for i in arc.getnames():
    buf = arc.extractfile(i).read()
    # openCVで画像を読込
    img = cv2.imdecode(np.frombuffer(buf, dtype=np.uint8), cv2.IMREAD_UNCHANGED)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    print(img.shape) # 例: (256, 256, 3)
    #
    # TODO: 画像処理など
    #
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    cv2.imwrite(f'output_{i}.png', img)

なお、出力された画像が元の tar に入れた画像と同じであることから、
変数 img にはRGBが反転している状態(imreadと同じ形式)で得られる.
そのため通常は cvtColor() でRGB反転を実施する必要がある.

with PyTorch Dataset

上記スクリプト例と同様に画像データを処理したいため、
Deep Learningのフレームワークで扱うことを考える.
PyTorchのDatasetで記述する場合は例えば以下の形で記述できる.

import cv2
import numpy as np
import albumentations as A
import tarfile
import torch
from torch.utils.data.dataset import Dataset

class SampleDataset(Dataset):
    def __init__(self, filename, augment):
        self.arc = tarfile.open(filename)
        self.files = [ i for i in self.arc.getnames() if self.arc.getmember(i).isfile() ]
        self.augment = augment

    def __len__(self):
        return len(self.files)

    def __getitem__(self, idx):
        buf = self.arc.extractfile(idx).read()
        img = cv2.imdecode(np.frombuffer(buf, dtype=np.uint8), cv2.IMREAD_UNCHANGED)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = self.augment(image=img)['image'].astype(np.float32)
        img = img.transpose(2, 0, 1)
        img = torch.from_numpy(img)
        return img

先と異なる点としてファイルの持ち方がある.
tarにはディレクトリ構造も保持でき、ディレクトリを含んだtarを読み込むと、
getnames() で取得する名前一覧にはディレクトリ自体も追加される.
具体例として testdir というディレクトリが存在する tar であれば下記のようになる.

import tarfile
tar = tarfile.open('test.tar')
files = tar.getnames()
print(files)
# ['sample1.jpg', 'sample2.jpg', 'testdir', 'testdir/sample3.jpg']

PyTorchのDataLoaderなどでの使用を考えるとファイル以外は取得したくない.
そのため getmember(name) で取得できる TarInfo の isfile() のプロパティを参照.
isfile() が True であるものだけを用いてリスト作成している.

import tarfile
tar = tarfile.open('test.tar')
files = [ f for f in tar.getnames() if tar.getmember(f).isfile() ]
print(files)
# ['sample1.jpg', 'sample2.jpg', 'testdir/sample3.jpg']

エラー処理やI/O実行効率の計測はできていないが、
ひとまず目的である tar のまま画像を取り扱うということは実現できる.

参考文献

• tarfile