Albumentations で XYZ軸回転のデータ拡張をするメモ

Albumentations で XYZ軸回転のデータ拡張をするメモです🦔
と言っても、Albumentations.Lambda()と自作のcv_rotate_3dを組み合わせただけの覚書です。

Google Colaboratory で動くサンプルノートブックを用意しているので、興味がある方は動かしてみてください。

---

以下、使い方メモ👀
サンプル画像は、ぱくたそ様の食後のケーキを利用しています。

cv_rotate_3dスクリプトダウンロード

!wget https://raw.githubusercontent.com/Kazuhito00/cv_rotate_3d/main/cv_rotate_3d.py

Albumentations準備

Lambda()内でcv_rotate_3d()を呼び出しています。
ランダム値の指定はrandom.uniform()を利用。

import random
import albumentations as album
from cv_rotate_3d import rotate_3d

def rotate_3d_for_albumentations(image, **params):
    theta = int(random.uniform(45, -45))
    phi = int(random.uniform(45, -45))
    gamma = int(random.uniform(45, -45))

    image = rotate_3d(
        image,
        theta=theta,       # X軸回転
        phi=phi,           # Y軸回転
        gamma=gamma,       # Z軸回転
        dx=0,              # X軸平行移動
        dy=0,              # Y軸平行移動
        dz=0,              # Z軸平行移動
        color=(0, 255, 0), # 余白色
    )
    return image

def preprocessing_augmentation_function(param_p=0.0):
    augmentation_function = album.Lambda(image=rotate_3d_for_albumentations, p=param_p)

    def augmentation(x):
        augmentation_image = augmentation_function(image=x)
        return augmentation_image['image']

    return augmentation

augmentation = preprocessing_augmentation_function(param_p=1.0) # param_pは適用確率

オリジナル画像

データ拡張結果