#!/usr/bin/env python
|
# -*- coding: utf-8 -*-
|
# @Time : 2020/10/26 16:25
|
# @Author : Scheaven
|
# @File : build.py
|
# @description:
|
|
import torchvision.transforms as T
|
|
from .transforms import *
|
from .autoaugment import AutoAugment
|
|
|
def build_transforms(cfg, is_train=True):
|
res = []
|
|
if is_train:
|
size_train = cfg.INPUT.SIZE_TRAIN
|
|
# augmix augmentation
|
do_augmix = cfg.INPUT.DO_AUGMIX
|
|
# auto augmentation
|
do_autoaug = cfg.INPUT.DO_AUTOAUG
|
total_iter = cfg.SOLVER.MAX_ITER
|
|
# horizontal filp
|
do_flip = cfg.INPUT.DO_FLIP
|
flip_prob = cfg.INPUT.FLIP_PROB
|
|
# padding
|
do_pad = cfg.INPUT.DO_PAD
|
padding = cfg.INPUT.PADDING
|
padding_mode = cfg.INPUT.PADDING_MODE
|
|
# color jitter
|
do_cj = cfg.INPUT.CJ.ENABLED
|
cj_prob = cfg.INPUT.CJ.PROB
|
cj_brightness = cfg.INPUT.CJ.BRIGHTNESS
|
cj_contrast = cfg.INPUT.CJ.CONTRAST
|
cj_saturation = cfg.INPUT.CJ.SATURATION
|
cj_hue = cfg.INPUT.CJ.HUE
|
|
# random erasing
|
do_rea = cfg.INPUT.REA.ENABLED
|
rea_prob = cfg.INPUT.REA.PROB
|
rea_mean = cfg.INPUT.REA.MEAN
|
# random patch
|
do_rpt = cfg.INPUT.RPT.ENABLED
|
rpt_prob = cfg.INPUT.RPT.PROB
|
|
if do_autoaug:
|
res.append(AutoAugment(total_iter))
|
res.append(T.Resize(size_train, interpolation=3))
|
if do_flip:
|
res.append(T.RandomHorizontalFlip(p=flip_prob))
|
if do_pad:
|
res.extend([T.Pad(padding, padding_mode=padding_mode),
|
T.RandomCrop(size_train)])
|
if do_cj:
|
res.append(T.RandomApply([T.ColorJitter(cj_brightness, cj_contrast, cj_saturation, cj_hue)], p=cj_prob))
|
if do_augmix:
|
res.append(AugMix())
|
if do_rea:
|
res.append(RandomErasing(probability=rea_prob, mean=rea_mean))
|
if do_rpt:
|
res.append(RandomPatch(prob_happen=rpt_prob))
|
else:
|
size_test = cfg.INPUT.SIZE_TEST
|
res.append(T.Resize(size_test, interpolation=3))
|
res.append(ToTensor())
|
res.append(T.Normalize(mean=[0.485*255, 0.456*255, 0.406*255], std=[0.229*255, 0.224*255, 0.225*255]))
|
return T.Compose(res)
|
