まえがき
pytorchでの学習では画像をnp.ndarrayなどの形式で読み込んだあとに，並進・回転・ガウスノイズ付加など，様々な変換をデータ拡張で行います．
その一つにJPEGで圧縮してブロックノイズを付加するというものがあります．データ拡張用のライブラリであるimgaugやalbumentationsでも実装されています．
これが，メモリ上でJPEG圧縮をやってくれているとばかり思っていたのに，一旦ディスクに保存するというコードになっていた（だから遅い），ということが発覚．
これはその顛末と対策です．
imgaugのJpegCompression
コードを見に行くと，PILでJPEGファイルとして一旦ディスクに書き込んだあと，imageioでそのファイルを読み込む，ということをしていました．
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を使用していました．
_, 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のコードはどうだろうと見に行くと，なんとやっぱり一旦ファイルに書き込んでからそのファイルを読み込んでいました．
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(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(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に変換
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で同様にするなら以下のコードです．
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ランタイムで計測すると似たような結果になってしまうのは謎．キャッシュに入る？）
%%timeit
img.save(
buffer,
format='JPEG',
quality=10)
img2 = Image.open(buffer)
%%timeit
Image.fromarray(img_array).save(
buffer,
format='JPEG',
quality=10)
img2 = Image.open(buffer)
img_array2 = np.array(img2)
%%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に変換する場合との計測比較は今後の課題です．