Pythonで画像をメモリ上でJPEG圧縮・展開する（たぶん）正しい方法

# まえがき

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

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

https://albumentations.ai/docs/api_reference/full_reference/#albumentations.augmentations.transforms.ImageCompression

https://imgaug.readthedocs.io/en/latest/source/overview/arithmetic.html?highlight=jpeg#jpegcompression

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

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

## imgaugのJpegCompression

[コード](https://github.com/aleju/imgaug/blob/0101108d4fed06bc5056c4a03e2bcb0216dac326/imgaug/augmenters/arithmetic.py#L1746)を見に行くと，PILでJPEGファイルとして一旦ディスクに書き込んだあと，imageioでそのファイルを読み込む，ということをしていました．

```python: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ならなおさら）．．．

https://docs.python.org/ja/3/library/tempfile.html#tempfile.NamedTemporaryFile


## albumentationsのImageCompression

[コード](https://github.com/albumentations-team/albumentations/blob/8428971dcb5de3f694c400c676cf5fb9029188bc/albumentations/augmentations/functional.py#L674)では，OpenCVの`imencode`と`imdecode`を使用していました．

```python: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圧縮・展開する方法として，あちこちで見かけます．

https://qiita.com/ka10ryu1/items/5fed6b4c8f29163d0d65

https://opqrstuvcut.github.io/blog/p/imencodeとimdecodeによるメモリ上での画像圧縮/


## OpenCVのimencode

では[OpenCVのimencode/imdecodeのコード](https://github.com/opencv/opencv/blob/2ebdc047875436f9e2dee31edb7f137a0df6c2cf/modules/imgcodecs/src/loadsave.cpp#L970)はどうだろうと見に行くと，なんとやっぱり一旦ファイルに書き込んでからそのファイルを読み込んでいました．

https://docs.opencv.org/3.0-beta/modules/imgcodecs/doc/reading_and_writing_images.html#imdecode

```cpp: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](https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.Image.save)

https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.Image.save

```text:使い方
PIL.Image.save(fp, format=None, **params)
fp – A filename (string), pathlib.Path object or file object.
```

つまりエンコードするには，`PIL.Image.save`の`fp`にfile objectを与えれば良さそうです．Pythonの[ファイルオブジェクト](https://docs.python.org/ja/3/c-api/file.html)にはバッファが使えるので，ファイル名を渡すのではなく，バッファ（メモリ領域）を渡すことにします．

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

- [PIL.Image.open](https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.open)

https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.open

```text:使い方
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に変換

```python: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で同様にするなら以下のコードです．

```python: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ランタイムで計測すると似たような結果になってしまうのは謎．キャッシュに入る？）


```python:1000 loops, best of 5: 373 µs per loop
%%timeit
img.save(
    buffer,
    format='JPEG',
    quality=10)
img2 = Image.open(buffer)
```


```python: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)
```


```python: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`が使えます．

https://pytorch.org/vision/main/generated/torchvision.io.decode_jpeg.html

```python:
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に変換する場合との計測比較は今後の課題です．