まえがき

pytorchでの学習では画像をnp.ndarrayなどの形式で読み込んだあとに，並進・回転・ガウスノイズ付加など，様々な変換をデータ拡張で行います．

その一つにJPEGで圧縮してブロックノイズを付加するというものがあります．データ拡張用のライブラリであるimgaugやalbumentationsでも実装されています．

これが，メモリ上でJPEG圧縮をやってくれているとばかり思っていたのに，一旦ディスクに保存するというコードになっていた（だから遅い），ということが発覚．

これはその顛末と対策です．

imgaugのJpegCompression

コードを見に行くと，PILでJPEGファイルとして一旦ディスクに書き込んだあと，imageioでそのファイルを読み込む，ということをしていました．

imgaug/imgaug/augmenters/arithmetic.py

    image_pil = PIL.Image.fromarray(image)
    with tempfile.NamedTemporaryFile(mode="wb+", suffix=".jpg") as f:
        image_pil.save(f, quality=quality)
        略
        image = imageio.imread(f, pilmode=pilmode, format="jpeg")
    略
    return image

dataloaderでもディスクのランダムアクセスが発生しているのに，その途中で（おそらく/tmp/以下に）tempfileを作っているのだから，実行時間がさらに遅くなる（HDDならなおさら）．．．

albumentationsのImageCompression

コードでは，OpenCVのimencodeとimdecodeを使用していました．

albumentations/albumentations/augmentations/functional.py

    _, encoded_img = cv2.imencode(image_type, img, (int(quality_flag), quality))
    img = cv2.imdecode(encoded_img, cv2.IMREAD_UNCHANGED)

この方法はメモリ上でJPEG圧縮・展開する方法として，あちこちで見かけます．

OpenCVのimencode

ではOpenCVのimencode/imdecodeのコードはどうだろうと見に行くと，なんとやっぱり一旦ファイルに書き込んでからそのファイルを読み込んでいました．

opencv/modules/imgcodecs/src/loadsave.cpp

        String filename = tempfile();
        code = encoder->setDestination(filename);
        略
        code = encoder->write(image, params);
        略
        FILE* f = fopen( filename.c_str(), "rb" );
        略
        buf.resize(fread( &buf[0], 1, buf.size(), f ));

encoder->writeの動作が追いきれていませんが，おそらくtempfileを作成しているのでしょう．

ここまでのまとめ

ということで，imgaugもalbumentationsもopencvも，JPEGの圧縮展開はメモリ上ではやっていません．

大量の画像に対してJPEG圧縮展開をメモリ上で効率的に行うには，どうしたら良いでしょう．

対策：PIL.Image

やり方は色々ありそうですが，PILを用いた方法を見つけました．

save&open

マニュアルによると，

PIL.Image.save

使い方

PIL.Image.save(fp, format=None, **params)
fp – A filename (string), pathlib.Path object or file object.

つまりエンコードするには，PIL.Image.saveのfpにfile objectを与えれば良さそうです．Pythonのファイルオブジェクトにはバッファが使えるので，ファイル名を渡すのではなく，バッファ（メモリ領域）を渡すことにします．

逆にデコードするには，PIL.Image.openのfpにバッファを渡せば良さそうです．

PIL.Image.open

使い方

PIL.Image.open(fp, mode='r', formats=None)
fp – A filename (string), pathlib.Path object or a file object. 略

コード

以下のコードの概要です．

ファイルimage.jpgをPIL.Image.openで読み込んだImageをndarrayに変換（単にndarrayを起点にしたいだけです）
PIL.Imageに変換してからsaveでbufferに書き込む．qualityパラメータは10．
PIL.Image.openでbufferからImageとして読み込み，ndarrayに変換

minimal code

from io import BytesIO
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt


img = Image.open('image.jpg')
img_array = np.array(img)
plt.imshow(img_array)

buffer = BytesIO()
Image.fromarray(img_array).save(
    buffer,
    format='JPEG',
    quality=10)

img2 = Image.open(buffer)
img_array2 = np.array(img2)
plt.imshow(img_array2)

OpenCVとの比較

opencvで同様にするなら以下のコードです．

minimal code of opencv


import cv2
img_cv = cv2.imread('image.jpg')

_, encoded_array = cv2.imencode(
    '.jpg',
    img_cv,
    (int(cv2.IMWRITE_JPEG_QUALITY), 10))

decoded_array = cv2.imdecode(
    encoded_array, 
    flags=cv2.IMREAD_UNCHANGED)

plt.imshow(decoded_array[:,:,::-1])  # BGR2RGB

速度比較

PIL.ImageでJPEG圧縮・展開をすると370µs = 0.37ms
ndarray-> PIL.ImageでJPEG圧縮・展開->ndarrayなら900µs = 0.9ms
opencvのimencode/imdecodeだと1.9ms

学習中ならデータローダーのディスクアクセスもあるためにimencode/imdecodeはもっと遅くなると予想します．

以下Google Colagoratoryのセルで実行時間計測した結果です．（ランタイムのタイプはNone=CPUで計測．GPUランタイムで計測すると似たような結果になってしまうのは謎．キャッシュに入る？）

1000 loops, best of 5: 373 µs per loop

%%timeit
img.save(
    buffer,
    format='JPEG',
    quality=10)
img2 = Image.open(buffer)

1000 loops, best of 5: 909 µs per loop

%%timeit
Image.fromarray(img_array).save(
    buffer,
    format='JPEG',
    quality=10)
img2 = Image.open(buffer)
img_array2 = np.array(img2)

1000 loops, best of 5: 1.91 ms per loop

%%timeit
_, encoded_array = cv2.imencode(
    '.jpg',
    img,
    (int(cv2.IMWRITE_JPEG_QUALITY), 10))
decoded_array = cv2.imdecode(
    encoded_array, 
    flags=cv2.IMREAD_UNCHANGED)

おまけ

データ拡張に使うのだから直接tensorにしてしまう方法があります．その場合にはtorchvisionのdecode_jpegが使えます．

import torch
from torch import as_tensor
from torchvision.io import decode_jpeg
from io import BytesIO
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt


img = Image.open('image.jpg')
img_array = np.array(img)
plt.imshow(img_array)

buffer = BytesIO()
Image.fromarray(img_array).save(
    buffer,
    format='JPEG',
    quality=10)


img_tensor = decode_jpeg(
    as_tensor(buffer.getbuffer(),
              dtype=torch.uint8))
plt.imshow(img_tensor.permute(1, 2, 0).numpy())

ただしas_tensorが遅いので，bufferがGPUメモリ上にある場合（データ転送が発生しない場合）に限って有利でしょう．Colab上では80ms程度とかなり低速でした．

PIL.ImageでJPEG圧縮・展開してからtensorに変換する場合との計測比較は今後の課題です．