TorchVision v2(version 0.16) について

以前から便利であったTorchVisionにおいてデータ拡張関連の部分がさらにアップデートされたようです．また実装に関しても，従来のライブラリにあったものは引き継がれているようなので，互換性があり移行は非常に楽です．

本記事では，以下のリリースノートを参考にGoogle Colabを使用しながらアップデート内容と実装方法に触れていきます．
https://github.com/pytorch/vision/releases/tag/v0.16.0

なお，0.17まではベータ版らしいので，次のアップデートではさらなる進化が期待できます．

セグメンテーションや物体検出，ビデオ等の3次元データに対応

ついにセグメンテーションや物体検出などに対するデータ拡張が公式で実装されました．
初めてセグメンテーションタスクでデータ拡張をするとき，なぜこれに対応していないのか疑問でした．しかも，データ拡張の関数の引数にSeed値がないため，自分でSeedを固定する必要があったりして少し面倒でした．

それではColabを使用して実際に実装してみましょう．

まず，Googleドライブのマウントを行います．

segmentation.ipynb

from google.colab import drive
drive.mount('/content/drive')
%cd ./drive/MyDrive

次に任意のフォルダに対象画像をアップロードし，パスを指定します．

segmentation.ipynb

# 適用画像のパスを指定
img_path = r"/content/drive/MyDrive/torchvision_v2/27861043_s.jpg"
mask_path = r"/content/drive/MyDrive/torchvision_v2/27861043_s_mask.jpg"

画像を表示する関数を定義します．

segmentation.ipynb

import matplotlib.pyplot as plt

# 画像表示
def show_two_images(img, mask):
    # サブプロットの設定
    fig, axes = plt.subplots(1, 2, figsize=(10, 5))

    # 1つ目の画像の表示
    axes[0].imshow(img)
    axes[0].axis('off')  # 軸を非表示にする
    axes[0].set_title('Image')

    # 2つ目の画像の表示
    axes[1].imshow(mask)
    axes[1].axis('off')  # 軸を非表示にする
    axes[1].set_title('Mask')

    # グラフを表示
    plt.show()

ここで，画像拡張を行い，オリジナル画像と拡張後の画像を比較します．

segmentation.ipynb

# セグメンテーション向けデータ拡張
from torchvision.transforms import v2
from PIL import Image

img = Image.open(img_path)
mask = Image.open(mask_path)

tf = v2.Compose([
                v2.RandomAffine(
                    degrees=(-40, 40), translate=(0.30, 0.15), scale=(0.8, 1.2)
                ),
                v2.RandomHorizontalFlip(0.5),
                ])
tf_img, tf_mask = tf(img, mask)

show_two_images(img, mask)
show_two_images(tf_img, tf_mask)

適用結果は以下のようになります．

マスク画像と入力画像に対して全く同じ処理ができていますね．
処理自体も引数にマスク画像を追加するだけで行ってくれるため，特にマスク画像であることを明示しなくても処理してくれるのもいいですね．

次に3次元画像に適用してみましょう．
ここで３次元画像は同じ画像を重ねた画像を3次元画像とします．

3Dimage.ipynb

import matplotlib.pyplot as plt
def show_3d_image(img):
    # サブプロットの設定
    fig, axes = plt.subplots(1, 5, figsize=(10, 5))

    for i in range(5):
      axes[i].imshow(img[i])
      axes[i].axis('off')  # 軸を非表示にする
      axes[i].set_title(f'Image {i}')

    # グラフを表示
    plt.show()

# 同じ画像を５枚重ねることで擬似的な三次元画像を作成
def make_3d_image(img):
  image_3d = []
  for _ in range(5):
    image_3d.append(img)

  return image_3d

3Dimage.ipynb

# 三次元データ用データ拡張
from torchvision.transforms import v2
from PIL import Image

img = Image.open(img_path)

tf = v2.Compose([
                v2.RandomAffine(
                    degrees=(-40, 40), translate=(0.30, 0.15), scale=(0.8, 1.2)
                ),
                v2.RandomHorizontalFlip(0.5),
                ])

image_3d = make_3d_image(img)
tf_image_3d = tf(image_3d)

show_3d_image(image_3d)
show_3d_image(tf_image_3d)

適用結果を示します．

すべての画像に対して同じ処理ができていますね．
3次元画像の型をlistで行ったんですが，普通に実行できましたね．前から出来てたっけ？

データ拡張手法のMixUpやCutMixに対応

これを使いたくて調べてたらPytorchの公式実装に追加されていることを知り，本記事を書くきっかけになりました．
コードは公式が実装例をアップロードしているので，これを参考に実行してみてください．
https://pytorch.org/vision/master/auto_examples/transforms/plot_cutmix_mixup.html
データローダーで読み込んだあとに，一行足すだけで良いのでかなり楽ですね．

また，Video等の三次元データにも対応しているので（型をtv_tensors.Videoにする必要はある)かなり便利ですね．

データ拡張の実行時間の高速化

リリースノートによると10%~40%の高速化に成功しているようです．

気が向いたら今後実装して計測してみます．