【Python】データ拡張ライブラリAlbumentationsの設定保存・復元

ディープラーニングで実験するときにどんなデータ拡張を利用したのかファイルとして保存しておきたかったのですが、 データ拡張ライブラリAlbumentationsの最新版には既にその機能があったのでメモします。

from albumentations import Compose
from albumentations.augmentations.transforms import Resize, HorizontalFlip, RandomSizedCrop, HueSaturationValue
from albumentations.core.serialization import save, load

def get_compose(crop_min_max, image_height, image_width, hue_shift, saturation_shift, value_shift):
    return Compose([Resize(image_height, image_width, p=1.0),
                     HorizontalFlip(p=0.5),
                     RandomSizedCrop(crop_min_max, image_height, image_width, p=1.0),
                     HueSaturationValue(hue_shift, saturation_shift, value_shift, p=1.0)])

# Resize image to 256x256
image_size = 256
# Crop 80 - 100% of image
crop_min = image_size*80//100
crop_max = image_size
crop_min_max = (crop_min, crop_max)
# HSV shift limits
hue_shift = 10
saturation_shift = 10
value_shift = 10
# Get compose
compose = get_compose(crop_min_max, image_size, image_size, hue_shift, saturation_shift, value_shift)

# Save compose in yaml format
save(compose, 'data_augmentations.yaml', data_format='yaml')
# Load compose from file
compose = load('data_augmentations.yaml', data_format='yaml')
print(compose)

data_augmentations.yamlにはこんな感じで保存されています。 一応ファイルを人目で見ても可読かと思います。

__version__: 0.4.3
transform:
  __class_fullname__: albumentations.core.composition.Compose
  additional_targets: {}
  bbox_params: null
  keypoint_params: null
  p: 1.0
  transforms:
  - __class_fullname__: albumentations.augmentations.transforms.Resize
    always_apply: false
    height: 256
    interpolation: 1
    p: 1.0
    width: 256
  - __class_fullname__: albumentations.augmentations.transforms.HorizontalFlip
    always_apply: false
    p: 0.5
  - __class_fullname__: albumentations.augmentations.transforms.RandomSizedCrop
    always_apply: false
    height: 256
    interpolation: 1
    min_max_height:
    - 204
    - 256
    p: 1.0
    w2h_ratio: 1.0
    width: 256
  - __class_fullname__: albumentations.augmentations.transforms.HueSaturationValue
    always_apply: false
    hue_shift_limit:
    - -10
    - 10
    p: 1.0
    sat_shift_limit:
    - -10
    - 10
    val_shift_limit:
    - -10
    - 10