def make_augmented(images, targets):
aug_params = [
OneOf([
Blur(blur_limit=5, p=1.),
RandomGamma(gamma_limit=(50, 150), p=1.),
HueSaturationValue(hue_shift_limit=20,
sat_shift_limit=30,
val_shift_limit=20,
p=1.),
RGBShift(r_shift_limit=15, g_shift_limit=5, b_shift_limit=15,
p=1.),
RandomBrightness(limit=.25, p=1.),
RandomContrast(limit=.25, p=1.),
MedianBlur(blur_limit=5, p=1.),
CLAHE(clip_limit=2.0, tile_grid_size=(8, 8), p=1.)
],
p=1.),
HorizontalFlip(p=1),
GaussNoise(var_limit=(10.0, 120.0), p=0.4)
]
labels = targets[:, 1]
bboxes = targets[:, 2:]
bboxes_new = torch.zeros((len(bboxes), 4))
bboxes_new[:, 0] = (bboxes[:, 0] - bboxes[:, 2] / 2)
bboxes_new[:, 1] = (bboxes[:, 1] - bboxes[:, 3] / 2)
bboxes_new[:, 2] = (bboxes[:, 0] + bboxes[:, 2] / 2)
bboxes_new[:, 3] = (bboxes[:, 1] + bboxes[:, 3] / 2)
# bboxes_new = bboxes_new.numpy()
annotations = {
'image': images,
'bboxes': bboxes_new.numpy(),
'labels': labels.numpy()
}
aug = get_aug(aug_params)
augmented = aug(**annotations)
augmented_image = augmented['image']
augmented_bboxes = np.array(augmented['bboxes'])
augmented_labels = np.array(augmented['labels'])
bboxes[:, 0] = torch.from_numpy(
((augmented_bboxes[:, 0] + augmented_bboxes[:, 2]) / 2))
bboxes[:, 1] = torch.from_numpy(
((augmented_bboxes[:, 1] + augmented_bboxes[:, 3]) / 2))
bboxes[:, 2] = torch.from_numpy(
(augmented_bboxes[:, 2] - augmented_bboxes[:, 0]))
bboxes[:, 3] = torch.from_numpy(
(augmented_bboxes[:, 3] - augmented_bboxes[:, 1]))
images = torch.from_numpy(augmented_image).float()
targets[:, 1] = torch.from_numpy(augmented_labels).float()
targets[:, 2:] = bboxes
# visualize(annotations, 'test_before.jpg')
# time.sleep(1)
# visualize(augmented, 'test_after.jpg')
# time.sleep(10)
return images, targets
def strong_aug_pixel(p=.5):
print('[DATA]: strong aug pixel')
from albumentations import (
# HorizontalFlip, IAAPerspective, ShiftScaleRotate, CLAHE, RandomRotate90,
Transpose,
ShiftScaleRotate,
Blur,
OpticalDistortion,
GridDistortion,
HueSaturationValue,
MultiplicativeNoise,
IAAAdditiveGaussianNoise,
GaussNoise,
MotionBlur,
MedianBlur,
RandomBrightnessContrast,
IAAPiecewiseAffine,
IAASharpen,
IAAEmboss,
Flip,
OneOf,
Compose,
JpegCompression,
CLAHE)
return Compose(
[
# RandomRotate90(),
# Flip(),
# Transpose(),
OneOf([
MultiplicativeNoise(multiplier=[0.5, 1.5], per_channel=True),
JpegCompression(quality_lower=39, quality_upper=80)
],
p=0.2),
OneOf([
IAAAdditiveGaussianNoise(),
GaussNoise(),
], p=0.2),
OneOf([
MotionBlur(p=.2),
MedianBlur(blur_limit=3, p=0.1),
Blur(blur_limit=3, p=0.1),
],
p=0.2),
# ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.2, rotate_limit=45, p=0.2),
# OneOf([
# OpticalDistortion(p=0.3),
# GridDistortion(p=.1),
# IAAPiecewiseAffine(p=0.3),
# ], p=0.2),
OneOf([
CLAHE(clip_limit=2),
IAASharpen(),
IAAEmboss(),
RandomBrightnessContrast(),
],
p=0.3),
HueSaturationValue(p=0.3),
],
p=p)
def _get_data_loader(imgs, trn_df, vld_df):
import albumentations as A
from albumentations import (
Rotate, HorizontalFlip, IAAPerspective, ShiftScaleRotate, CLAHE,
RandomRotate90, Transpose, ShiftScaleRotate, Blur, OpticalDistortion,
GridDistortion, HueSaturationValue, IAAAdditiveGaussianNoise,
GaussNoise, MotionBlur, MedianBlur, RandomBrightnessContrast,
IAAPiecewiseAffine, IAASharpen, IAAEmboss, Flip, OneOf, Compose)
from albumentations.pytorch import ToTensor, ToTensorV2
train_transforms = A.Compose([
Rotate(20),
OneOf([
IAAAdditiveGaussianNoise(),
GaussNoise(),
], p=0.2),
OneOf([
MotionBlur(p=.2),
MedianBlur(blur_limit=3, p=0.1),
Blur(blur_limit=3, p=0.1),
],
p=0.2),
ShiftScaleRotate(
shift_limit=0.0625, scale_limit=0.2, rotate_limit=45, p=0.2),
OneOf([
OpticalDistortion(p=0.3),
GridDistortion(p=.1),
IAAPiecewiseAffine(p=0.3),
],
p=0.2),
OneOf([
CLAHE(clip_limit=2),
IAASharpen(),
IAAEmboss(),
RandomBrightnessContrast(),
],
p=0.3),
HueSaturationValue(p=0.3),
ToTensor()
],
p=1.0)
valid_transforms = A.Compose([ToTensor()])
from torch.utils.data import Dataset, DataLoader
trn_dataset = BangaliDataset(imgs=imgs,
label_df=trn_df,
transform=train_transforms)
vld_dataset = BangaliDataset(imgs=imgs,
label_df=vld_df,
transform=valid_transforms)
trn_loader = DataLoader(trn_dataset,
shuffle=True,
num_workers=NUM_WORKERS,
batch_size=BATCH_SIZE)
vld_loader = DataLoader(vld_dataset,
shuffle=False,
num_workers=NUM_WORKERS,
batch_size=BATCH_SIZE)
return trn_loader, vld_loader
import numpy as np
import cv2
from plasma.training.data import Dataset, augmentations as augs
from albumentations import HorizontalFlip, OneOf, RandomBrightnessContrast, RandomGamma, ShiftScaleRotate, Compose, \
CoarseDropout
aug = Compose([
CoarseDropout(min_holes=0, max_holes=3, max_height=32, max_width=32,
p=0.8),
augs.MinEdgeCrop(always_apply=True),
HorizontalFlip(),
OneOf([
RandomGamma(),
RandomBrightnessContrast(),
], p=0.8),
ShiftScaleRotate(shift_limit=0.1,
rotate_limit=35,
scale_limit=0.2,
p=0.8,
border_mode=cv2.BORDER_CONSTANT),
])
class Data(Dataset):
def __init__(self, df, image_path):
super().__init__()
self.df = df.copy().reset_index(drop=True)
self.image_path = image_path
from albumentations import Compose, Resize, RandomCrop, Flip, HorizontalFlip, VerticalFlip, Transpose, RandomRotate90, \
ShiftScaleRotate, OneOf, Blur, MotionBlur, MedianBlur, GaussianBlur
from albumentations.pytorch import ToTensor
train_aug = Compose([
RandomCrop(height=96, width=96, p=0.2),
OneOf([
VerticalFlip(p=0.2),
HorizontalFlip(p=0.3),
Transpose(p=0.2),
RandomRotate90(p=0.2),
],
p=0.3),
ShiftScaleRotate(p=0.2),
OneOf([
Blur(p=0.2),
MotionBlur(p=0.2),
MedianBlur(p=0.2),
GaussianBlur(p=0.2),
],
p=0.3),
Resize(196, 196, always_apply=True),
ToTensor()
])
valid_aug = Compose([Resize(196, 196, always_apply=True), ToTensor()])
def __init__(self,
imgs: Sequence[str] = None,
suffix: str = '.path',
line_width: int = 4,
im_transforms: Callable[[Any],
torch.Tensor] = transforms.Compose(
[]),
mode: str = 'path',
augmentation: bool = False,
valid_baselines: Sequence[str] = None,
merge_baselines: Dict[str, Sequence[str]] = None,
valid_regions: Sequence[str] = None,
merge_regions: Dict[str, Sequence[str]] = None):
"""
Reads a list of image-json pairs and creates a data set.
Args:
imgs (list):
suffix (int): Suffix to attach to image base name to load JSON
files from.
line_width (int): Height of the baseline in the scaled input.
target_size (tuple): Target size of the image as a (height, width) tuple.
mode (str): Either path, alto, page, xml, or None. In alto, page,
and xml mode the baseline paths and image data is
retrieved from an ALTO/PageXML file. In `None` mode
data is iteratively added through the `add` method.
augmentation (bool): Enable/disable augmentation.
valid_baselines (list): Sequence of valid baseline identifiers. If
`None` all are valid.
merge_baselines (dict): Sequence of baseline identifiers to merge.
Note that merging occurs after entities not
in valid_* have been discarded.
valid_regions (list): Sequence of valid region identifiers. If
`None` all are valid.
merge_regions (dict): Sequence of region identifiers to merge.
Note that merging occurs after entities not
in valid_* have been discarded.
"""
super().__init__()
self.mode = mode
self.im_mode = '1'
self.aug = None
self.targets = []
# n-th entry contains semantic of n-th class
self.class_mapping = {
'aux': {
'_start_separator': 0,
'_end_separator': 1
},
'baselines': {},
'regions': {}
}
self.num_classes = 2
self.mbl_dict = merge_baselines if merge_baselines is not None else {}
self.mreg_dict = merge_regions if merge_regions is not None else {}
self.valid_baselines = valid_baselines
self.valid_regions = valid_regions
if mode in ['alto', 'page', 'xml']:
if mode == 'alto':
fn = parse_alto
elif mode == 'page':
fn = parse_page
elif mode == 'xml':
fn = parse_xml
im_paths = []
self.targets = []
for img in imgs:
try:
data = fn(img)
im_paths.append(data['image'])
lines = defaultdict(list)
for line in data['lines']:
if valid_baselines is None or line[
'script'] in valid_baselines:
lines[self.mbl_dict.get(line['script'],
line['script'])].append(
line['baseline'])
regions = defaultdict(list)
for k, v in data['regions'].items():
if valid_regions is None or k in valid_regions:
regions[self.mreg_dict.get(k, k)].extend(v)
data['regions'] = regions
self.targets.append({
'baselines': lines,
'regions': data['regions']
})
except KrakenInputException as e:
logger.warning(e)
continue
# get line types
imgs = im_paths
# calculate class mapping
line_types = set()
region_types = set()
for page in self.targets:
for line_type in page['baselines'].keys():
line_types.add(line_type)
for reg_type in page['regions'].keys():
region_types.add(reg_type)
idx = -1
for idx, line_type in enumerate(line_types):
self.class_mapping['baselines'][
line_type] = idx + self.num_classes
self.num_classes += idx + 1
idx = -1
for idx, reg_type in enumerate(region_types):
self.class_mapping['regions'][
reg_type] = idx + self.num_classes
self.num_classes += idx + 1
elif mode == 'path':
pass
elif mode is None:
imgs = []
else:
raise Exception('invalid dataset mode')
if augmentation:
from albumentations import (
Compose,
ToFloat,
FromFloat,
RandomRotate90,
Flip,
OneOf,
MotionBlur,
MedianBlur,
Blur,
ShiftScaleRotate,
OpticalDistortion,
ElasticTransform,
RandomBrightnessContrast,
HueSaturationValue,
)
self.aug = Compose([
ToFloat(),
RandomRotate90(),
Flip(),
OneOf([
MotionBlur(p=0.2),
MedianBlur(blur_limit=3, p=0.1),
Blur(blur_limit=3, p=0.1),
],
p=0.2),
ShiftScaleRotate(shift_limit=0.0625,
scale_limit=0.2,
rotate_limit=45,
p=0.2),
OneOf([
OpticalDistortion(p=0.3),
ElasticTransform(p=0.1),
],
p=0.2),
HueSaturationValue(hue_shift_limit=20,
sat_shift_limit=0.1,
val_shift_limit=0.1,
p=0.3),
],
p=0.5)
self.imgs = imgs
self.line_width = line_width
# split image transforms into two. one part giving the final PIL image
# before conversion to a tensor and the actual tensor conversion part.
self.head_transforms = transforms.Compose(im_transforms.transforms[:2])
self.tail_transforms = transforms.Compose(im_transforms.transforms[2:])
self.seg_type = None
def main() :
# mode argument
args = argparse.ArgumentParser()
args.add_argument("--letter", type=str, default=" ,.()\'\"?!01234567890abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ가각간갇갈갉갊감갑값갓갔강갖갗같갚갛개객갠갤갬갭갯갰갱갸갹갼걀걋걍걔걘걜거걱건걷걸걺검겁것겄겅겆겉겊겋게겐겔겜겝겟겠겡겨격겪견겯결겸겹겻겼경곁계곈곌곕곗고곡곤곧골곪곬곯곰곱곳공곶과곽관괄괆괌괍괏광괘괜괠괩괬괭괴괵괸괼굄굅굇굉교굔굘굡굣구국군굳굴굵굶굻굼굽굿궁궂궈궉권궐궜궝궤궷귀귁귄귈귐귑귓규균귤그극근귿글긁금급긋긍긔기긱긴긷길긺김깁깃깅깆깊까깍깎깐깔깖깜깝깟깠깡깥깨깩깬깰깸깹깻깼깽꺄꺅꺌꺼꺽꺾껀껄껌껍껏껐껑께껙껜껨껫껭껴껸껼꼇꼈꼍꼐꼬꼭꼰꼲꼴꼼꼽꼿꽁꽂꽃꽈꽉꽐꽜꽝꽤꽥꽹꾀꾄꾈꾐꾑꾕꾜꾸꾹꾼꿀꿇꿈꿉꿋꿍꿎꿔꿜꿨꿩꿰꿱꿴꿸뀀뀁뀄뀌뀐뀔뀜뀝뀨끄끅끈끊끌끎끓끔끕끗끙끝끼끽낀낄낌낍낏낑나낙낚난낟날낡낢남납낫났낭낮낯낱낳내낵낸낼냄냅냇냈냉냐냑냔냘냠냥너넉넋넌널넒넓넘넙넛넜넝넣네넥넨넬넴넵넷넸넹녀녁년녈념녑녔녕녘녜녠노녹논놀놂놈놉놋농높놓놔놘놜놨뇌뇐뇔뇜뇝뇟뇨뇩뇬뇰뇹뇻뇽누눅눈눋눌눔눕눗눙눠눴눼뉘뉜뉠뉨뉩뉴뉵뉼늄늅늉느늑는늘늙늚늠늡늣능늦늪늬늰늴니닉닌닐닒님닙닛닝닢다닥닦단닫달닭닮닯닳담답닷닸당닺닻닿대댁댄댈댐댑댓댔댕댜더덕덖던덛덜덞덟덤덥덧덩덫덮데덱덴델뎀뎁뎃뎄뎅뎌뎐뎔뎠뎡뎨뎬도독돈돋돌돎돐돔돕돗동돛돝돠돤돨돼됐되된될됨됩됫됴두둑둔둘둠둡둣둥둬뒀뒈뒝뒤뒨뒬뒵뒷뒹듀듄듈듐듕드득든듣들듦듬듭듯등듸디딕딘딛딜딤딥딧딨딩딪따딱딴딸땀땁땃땄땅땋때땍땐땔땜땝땟땠땡떠떡떤떨떪떫떰떱떳떴떵떻떼떽뗀뗄뗌뗍뗏뗐뗑뗘뗬또똑똔똘똥똬똴뙈뙤뙨뚜뚝뚠뚤뚫뚬뚱뛔뛰뛴뛸뜀뜁뜅뜨뜩뜬뜯뜰뜸뜹뜻띄띈띌띔띕띠띤띨띰띱띳띵라락란랄람랍랏랐랑랒랖랗래랙랜랠램랩랫랬랭랴략랸럇량러럭런럴럼럽럿렀렁렇레렉렌렐렘렙렛렝려력련렬렴렵렷렸령례롄롑롓로록론롤롬롭롯롱롸롼뢍뢨뢰뢴뢸룀룁룃룅료룐룔룝룟룡루룩룬룰룸룹룻룽뤄뤘뤠뤼뤽륀륄륌륏륑류륙륜률륨륩륫륭르륵른를름릅릇릉릊릍릎리릭린릴림립릿링마막만많맏말맑맒맘맙맛망맞맡맣매맥맨맬맴맵맷맸맹맺먀먁먈먕머먹먼멀멂멈멉멋멍멎멓메멕멘멜멤멥멧멨멩며멱면멸몃몄명몇몌모목몫몬몰몲몸몹못몽뫄뫈뫘뫙뫼묀묄묍묏묑묘묜묠묩묫무묵묶문묻물묽묾뭄뭅뭇뭉뭍뭏뭐뭔뭘뭡뭣뭬뮈뮌뮐뮤뮨뮬뮴뮷므믄믈믐믓미믹민믿밀밂밈밉밋밌밍및밑바박밖밗반받발밝밞밟밤밥밧방밭배백밴밸뱀뱁뱃뱄뱅뱉뱌뱍뱐뱝버벅번벋벌벎범법벗벙벚베벡벤벧벨벰벱벳벴벵벼벽변별볍볏볐병볕볘볜보복볶본볼봄봅봇봉봐봔봤봬뵀뵈뵉뵌뵐뵘뵙뵤뵨부북분붇불붉붊붐붑붓붕붙붚붜붤붰붸뷔뷕뷘뷜뷩뷰뷴뷸븀븃븅브븍븐블븜븝븟비빅빈빌빎빔빕빗빙빚빛빠빡빤빨빪빰빱빳빴빵빻빼빽뺀뺄뺌뺍뺏뺐뺑뺘뺙뺨뻐뻑뻔뻗뻘뻠뻣뻤뻥뻬뼁뼈뼉뼘뼙뼛뼜뼝뽀뽁뽄뽈뽐뽑뽕뾔뾰뿅뿌뿍뿐뿔뿜뿟뿡쀼쁑쁘쁜쁠쁨쁩삐삑삔삘삠삡삣삥사삭삯산삳살삵삶삼삽삿샀상샅새색샌샐샘샙샛샜생샤샥샨샬샴샵샷샹섀섄섈섐섕서석섞섟선섣설섦섧섬섭섯섰성섶세섹센셀셈셉셋셌셍셔셕션셜셤셥셧셨셩셰셴셸솅소속솎손솔솖솜솝솟송솥솨솩솬솰솽쇄쇈쇌쇔쇗쇘쇠쇤쇨쇰쇱쇳쇼쇽숀숄숌숍숏숑수숙순숟술숨숩숫숭숯숱숲숴쉈쉐쉑쉔쉘쉠쉥쉬쉭쉰쉴쉼쉽쉿슁슈슉슐슘슛슝스슥슨슬슭슴습슷승시식신싣실싫심십싯싱싶싸싹싻싼쌀쌈쌉쌌쌍쌓쌔쌕쌘쌜쌤쌥쌨쌩썅써썩썬썰썲썸썹썼썽쎄쎈쎌쏀쏘쏙쏜쏟쏠쏢쏨쏩쏭쏴쏵쏸쐈쐐쐤쐬쐰쐴쐼쐽쑈쑤쑥쑨쑬쑴쑵쑹쒀쒔쒜쒸쒼쓩쓰쓱쓴쓸쓺쓿씀씁씌씐씔씜씨씩씬씰씸씹씻씽아악안앉않알앍앎앓암압앗았앙앝앞애액앤앨앰앱앳앴앵야약얀얄얇얌얍얏양얕얗얘얜얠얩어억언얹얻얼얽얾엄업없엇었엉엊엌엎에엑엔엘엠엡엣엥여역엮연열엶엷염엽엾엿였영옅옆옇예옌옐옘옙옛옜오옥온올옭옮옰옳옴옵옷옹옻와왁완왈왐왑왓왔왕왜왝왠왬왯왱외왹왼욀욈욉욋욍요욕욘욜욤욥욧용우욱운울욹욺움웁웃웅워웍원월웜웝웠웡웨웩웬웰웸웹웽위윅윈윌윔윕윗윙유육윤율윰윱윳융윷으윽은을읊음읍읏응읒읓읔읕읖읗의읩읜읠읨읫이익인일읽읾잃임입잇있잉잊잎자작잔잖잗잘잚잠잡잣잤장잦재잭잰잴잼잽잿쟀쟁쟈쟉쟌쟎쟐쟘쟝쟤쟨쟬저적전절젊점접젓정젖제젝젠젤젬젭젯젱져젼졀졈졉졌졍졔조족존졸졺좀좁좃종좆좇좋좌좍좔좝좟좡좨좼좽죄죈죌죔죕죗죙죠죡죤죵주죽준줄줅줆줌줍줏중줘줬줴쥐쥑쥔쥘쥠쥡쥣쥬쥰쥴쥼즈즉즌즐즘즙즛증지직진짇질짊짐집짓징짖짙짚짜짝짠짢짤짧짬짭짯짰짱째짹짼쨀쨈쨉쨋쨌쨍쨔쨘쨩쩌쩍쩐쩔쩜쩝쩟쩠쩡쩨쩽쪄쪘쪼쪽쫀쫄쫌쫍쫏쫑쫓쫘쫙쫠쫬쫴쬈쬐쬔쬘쬠쬡쭁쭈쭉쭌쭐쭘쭙쭝쭤쭸쭹쮜쮸쯔쯤쯧쯩찌찍찐찔찜찝찡찢찧차착찬찮찰참찹찻찼창찾채책챈챌챔챕챗챘챙챠챤챦챨챰챵처척천철첨첩첫첬청체첵첸첼쳄쳅쳇쳉쳐쳔쳤쳬쳰촁초촉촌촐촘촙촛총촤촨촬촹최쵠쵤쵬쵭쵯쵱쵸춈추축춘출춤춥춧충춰췄췌췐취췬췰췸췹췻췽츄츈츌츔츙츠측츤츨츰츱츳층치칙친칟칠칡침칩칫칭카칵칸칼캄캅캇캉캐캑캔캘캠캡캣캤캥캬캭컁커컥컨컫컬컴컵컷컸컹케켁켄켈켐켑켓켕켜켠켤켬켭켯켰켱켸코콕콘콜콤콥콧콩콰콱콴콸쾀쾅쾌쾡쾨쾰쿄쿠쿡쿤쿨쿰쿱쿳쿵쿼퀀퀄퀑퀘퀭퀴퀵퀸퀼큄큅큇큉큐큔큘큠크큭큰클큼큽킁키킥킨킬킴킵킷킹타탁탄탈탉탐탑탓탔탕태택탠탤탬탭탯탰탱탸턍터턱턴털턺텀텁텃텄텅테텍텐텔템텝텟텡텨텬텼톄톈토톡톤톨톰톱톳통톺톼퇀퇘퇴퇸툇툉툐투툭툰툴툼툽툿퉁퉈퉜퉤튀튁튄튈튐튑튕튜튠튤튬튱트특튼튿틀틂틈틉틋틔틘틜틤틥티틱틴틸팀팁팃팅파팍팎판팔팖팜팝팟팠팡팥패팩팬팰팸팹팻팼팽퍄퍅퍼퍽펀펄펌펍펏펐펑페펙펜펠펨펩펫펭펴편펼폄폅폈평폐폘폡폣포폭폰폴폼폽폿퐁퐈퐝푀푄표푠푤푭푯푸푹푼푿풀풂품풉풋풍풔풩퓌퓐퓔퓜퓟퓨퓬퓰퓸퓻퓽프픈플픔픕픗피픽핀필핌핍핏핑하학한할핥함합핫항해핵핸핼햄햅햇했행햐향허헉헌헐헒험헙헛헝헤헥헨헬헴헵헷헹혀혁현혈혐협혓혔형혜혠혤혭호혹혼홀홅홈홉홋홍홑화확환활홧황홰홱홴횃횅회획횐횔횝횟횡효횬횰횹횻후훅훈훌훑훔훗훙훠훤훨훰훵훼훽휀휄휑휘휙휜휠휨휩휫휭휴휵휸휼흄흇흉흐흑흔흖흗흘흙흠흡흣흥흩희흰흴흼흽힁히힉힌힐힘힙힛힝")
args.add_argument("--cuda", type=bool, default=True)
args.add_argument("--num_epochs", type=int, default=500)
args.add_argument("--mode", type=str, default="train")
args.add_argument("--print_iter", type=int, default=1000)
args.add_argument("--imgH", type=int, default=32)
args.add_argument("--imgW", type=int, default=200)
args.add_argument("--batch_size", type=int, default=192)
args.add_argument("--batch_max_length", type=int, default=45)
args.add_argument("--output_channel", type=int, default=320)
args.add_argument("--hidden_size", type=int, default=160)
args.add_argument("--input_channel", type=int, default=1)
args.add_argument("--lr", type=float, default=1.0)
args.add_argument("--rho", type=float, default=0.95)
args.add_argument("--eps", type=float, default=1e-8)
args.add_argument("--grad_clip", type=float, default=5)
args.add_argument("--PAD", type=bool, default=True)
args.add_argument("--data_filtering_off", type=bool, default=True)
args.add_argument("--num_fiducial", type=int, default=20)
args.add_argument('--extract', type=str, default='RCNN')
args.add_argument('--grid_mask', type=bool, default=True)
args.add_argument("--load_model", type=str, default='')
args.add_argument("--save_dir", type=str, default= '')
###########################################
opt = args.parse_args()
if not opt.save_dir:
opt.save_dir = f'{SaveDir_maker()}'
letter = opt.letter
lr = opt.lr
cuda = opt.cuda
num_epochs = opt.num_epochs
load_model = opt.load_model
batch = opt.batch_size
mode = opt.mode
prediction_dir = f'/tf/notebooks/models/{opt.load_model.split("/")[0]}/{opt.load_model.split("/")[1]}/prediction'
print_iter = opt.print_iter
imgH = opt.imgH
imgW = opt.imgW
converter = utils.AttnLabelConverter(opt.letter)
opt.num_classes = len(converter.character)
#########################################
device = torch.device('cuda') if cuda else torch.device('cpu')
new_model = www_model.STR(opt, device)
if load_model:
load_pretrained_model(load_model, new_model)
images = torch.FloatTensor(batch, 1, imgH, imgW)
texts = torch.IntTensor(batch * 1000)
lengths = torch.IntTensor(batch)
images = Variable(images)
texts = Variable(texts)
lengths = Variable(lengths)
#check parameter of model
print("------------------------------------------------------------")
total_params = sum(p.numel() for p in new_model.parameters())
print("num of parameter : ",total_params)
trainable_params = sum(p.numel() for p in new_model.parameters() if p.requires_grad)
print("num of trainable_ parameter :",trainable_params)
print("------------------------------------------------------------")
print('prediction directory : ' , prediction_dir)
if mode == 'train' :
print('data loading')
if opt.grid_mask :
transforms_train = Compose([dataloader_gridmask.GridMask(num_grid=(1,1), div_val=[10,2], p=0.3), OneOf([IAAAdditiveGaussianNoise(),
GaussNoise()], p=0.3 )])
train_loader = dataloader_gridmask.data_loader(DATASET_PATH, batch, imgH, imgW, transform = transforms_train, phase='train')
val_loader = data_loader(DATASET_PATH, 1, imgH, imgW, phase='val')
else:
train_loader = data_loader(
DATASET_PATH, batch, imgH, imgW, phase='train')
val_loader = data_loader(
DATASET_PATH, 1, imgH, imgW, phase='val')
params = [p for p in new_model.parameters() if p.requires_grad]
# optimizer = optim.Adam(params, lr=lr, betas=(0.5, 0.999))
optimizer = optim.Adadelta(params, lr= opt.lr, rho = opt.rho, eps = opt.eps)
lr_scheduler = optim.lr_scheduler.StepLR(optimizer,step_size=40,gamma=0.1)
#opt log
os.makedirs(f'{opt.save_dir}', exist_ok=True)
with open(f'{opt.save_dir}/opt.txt', 'a') as opt_file:
opt_log = '---------------------Options-----------------\n'
args = vars(opt)
for k, v in args.items():
opt_log +=f'{str(k)} : {str(v)}\n'
opt_log +='---------------------------------------------\n'
opt_file.write(opt_log)
print('train start')
train(num_epochs, new_model, device, train_loader, val_loader, images, texts, lengths, converter, optimizer, lr_scheduler, prediction_dir, print_iter, opt)
elif mode == 'test' :
print('test start')
test_loader = data_loader(
DATASET_PATH, 1, imgH, imgW, phase='test')
test(new_model, device, test_loader, images, texts, lengths, converter, prediction_dir, opt)
RandomRotate90, ElasticTransform, GridDistortion,
OpticalDistortion, OneOf, CLAHE, RandomContrast,
RandomGamma, RandomBrightness)
MODEL_NAME = 'linknet_heavy_aug'
original_height = 101
original_width = 101
aug = Compose([
VerticalFlip(p=0.5),
RandomRotate90(p=0.5),
OneOf([
ElasticTransform(
p=0.5, alpha=120, sigma=120 * 0.05, alpha_affine=120 * 0.03),
GridDistortion(p=0.5),
OpticalDistortion(p=1, distort_limit=2, shift_limit=0.5)
],
p=0.8),
CLAHE(p=0.8),
RandomContrast(p=0.8),
RandomBrightness(p=0.8),
RandomGamma(p=0.8)
])
if __name__ == '__main__':
directory = '/root/data/salt/'
n_fold = 8
depths = pd.read_csv(os.path.join("/root/data/salt/", 'depths.csv'))
depths.sort_values('z', inplace=True)
from albumentations import Compose, Resize, RandomCrop, Flip, HorizontalFlip, VerticalFlip, Transpose, RandomRotate90, \
ShiftScaleRotate, OneOf, OpticalDistortion, GridDistortion
from albumentations.pytorch import ToTensor
train_aug = Compose([
RandomCrop(height=96, width=96, p=0.2),
OneOf([
VerticalFlip(p=0.2),
HorizontalFlip(p=0.3),
Transpose(p=0.2),
RandomRotate90(p=0.2),
],
p=0.3),
ShiftScaleRotate(p=0.2),
OneOf([
OpticalDistortion(p=0.2),
GridDistortion(p=0.2),
], p=0.3),
Resize(128, 128, always_apply=True),
ToTensor()
])
valid_aug = Compose([Resize(128, 128, always_apply=True), ToTensor()])
#Augmentaions here are uising Albumentations library "https://github.com/albu/albumentations" developed by Kaggle Grandmasters
from albumentations import (PadIfNeeded, HorizontalFlip, VerticalFlip,
CenterCrop, Crop, Compose, Transpose,
RandomRotate90, ElasticTransform, GridDistortion,
OpticalDistortion, RandomSizedCrop, OneOf, CLAHE,
RandomBrightnessContrast, RandomGamma, RGBShift)
original_height = 224
original_width = 224
aug = Compose([
OneOf([
RandomSizedCrop(min_max_height=(50, 101),
height=original_height,
width=original_width,
p=0.5),
PadIfNeeded(
min_height=original_height, min_width=original_width, p=0.5)
],
p=1),
VerticalFlip(p=0.5),
RandomRotate90(p=0.5),
OneOf([
ElasticTransform(
p=0.4, alpha=120, sigma=120 * 0.05, alpha_affine=120 * 0.03),
GridDistortion(p=0.3),
OpticalDistortion(p=0.4, distort_limit=2, shift_limit=0.5)
],
p=0.8),
CLAHE(p=0.8),
RandomBrightnessContrast(p=0.8),
def __getitem__(self, idx):
file = self.files[idx]
if self.mode == 'train' or 'val':
file_name = 'train'
else:
file_name = 'test'
file_path = os.path.join(os.path.join(PATH, file_name + '_images'),
file)
image = cv2.imread(file_path)
mask = rleToMask(self.masks[idx])
mean = (0.485, 0.456, 0.406)
std = (0.229, 0.224, 0.225)
train_aug = Compose([
OneOf([
ShiftScaleRotate(),
VerticalFlip(p=0.8),
HorizontalFlip(p=0.8),
],
p=0.6),
OneOf([
RandomBrightnessContrast(),
MotionBlur(p=0.5),
MedianBlur(blur_limit=3, p=0.5),
Blur(blur_limit=3, p=0.5),
],
p=0.6),
OneOf([
IAAAdditiveGaussianNoise(),
GaussNoise(),
], p=0.6),
Normalize(mean=mean, std=std, p=1)
# ToTensor()
])
val_aug = Compose([Normalize(mean=mean, std=std, p=1), ToTensor()])
if self.mode == 'train':
crop_images, crop_masks, new_crop_images = [], [], []
crop_labels = np.zeros((num_crop, 4))
for i in range(num_crop):
crop_images.append(image[:, i * (1600 // num_crop):(i + 1) *
(1600 // num_crop), :])
crop_masks.append(mask[:, :, i * (1600 // num_crop):(i + 1) *
(1600 // num_crop)])
for i in range(num_crop):
for ch in range(4):
if (crop_masks[i][ch, :, :] == np.zeros(
(crop_masks[i][ch, :, :].shape))).all():
crop_labels[i][ch] = 0
else:
crop_labels[i][ch] = 1
for num in range(len(crop_images)):
if self.mode == 'train':
augmented = train_aug(image=crop_images[num])
new_crop_images.append(augmented['image'])
crop_labels[num] = np.array(crop_labels[num])
crop_labels[num] = torch.tensor(crop_labels[num],
dtype=torch.float32)
new_crop_images = np.transpose(np.array(new_crop_images),
(0, 3, 1, 2))
return new_crop_images, crop_labels
elif self.mode == 'val':
augmented = val_aug(image=image)
image = augmented['image']
label = np.array(self.labels[idx])
label = torch.tensor(label, dtype=torch.float32)
return (image, label)
def albumentations_transforms(
crop_size,
shorter_side,
low_scale,
high_scale,
img_mean,
img_std,
img_scale,
ignore_label,
num_stages,
dataset_type,
):
from albumentations import (
Normalize,
HorizontalFlip,
RandomCrop,
PadIfNeeded,
RandomScale,
LongestMaxSize,
SmallestMaxSize,
OneOf,
)
from albumentations.pytorch import ToTensorV2 as ToTensor
from densetorch.data import albumentations2densetorch
if dataset_type == "densetorch":
wrapper = albumentations2densetorch
elif dataset_type == "torchvision":
wrapper = albumentations2torchvision
else:
raise ValueError(f"Unknown dataset type: {dataset_type}")
common_transformations = [
Normalize(max_pixel_value=1.0 / img_scale, mean=img_mean, std=img_std),
ToTensor(),
]
train_transforms = []
for stage in range(num_stages):
train_transforms.append(
wrapper(
[
OneOf(
[
RandomScale(
scale_limit=(low_scale[stage], high_scale[stage])
),
LongestMaxSize(max_size=shorter_side[stage]),
SmallestMaxSize(max_size=shorter_side[stage]),
]
),
PadIfNeeded(
min_height=crop_size[stage],
min_width=crop_size[stage],
border_mode=cv2.BORDER_CONSTANT,
value=np.array(img_mean) / img_scale,
mask_value=ignore_label,
),
HorizontalFlip(p=0.5,),
RandomCrop(height=crop_size[stage], width=crop_size[stage],),
]
+ common_transformations
)
)
val_transforms = wrapper(common_transformations)
return train_transforms, val_transforms
from albumentations import Compose, ShiftScaleRotate, PadIfNeeded, RandomCrop, Resize, RandomSizedCrop, CLAHE, \
RandomRotate90, Flip, OneOf, MotionBlur, MedianBlur, Blur, CenterCrop, LongestMaxSize, HorizontalFlip, VerticalFlip, \
Transpose
from albumentations.pytorch import ToTensor
transform_train = Compose([
RandomRotate90(0.2),
Flip(p=0.2),
ShiftScaleRotate(),
OneOf([
MotionBlur(p=.2),
MedianBlur(blur_limit=3, p=0.1),
Blur(blur_limit=3, p=0.1),
],
p=0.2),
ToTensor()
])
transform_test = Compose([ToTensor()])
IMG_SIZE_RAW = 224
RAW_CROP_SIZE = 448
transform_train_raw = Compose([
Resize(RAW_CROP_SIZE, RAW_CROP_SIZE),
# CenterCrop(width=IMG_SIZE_RAW, height=IMG_SIZE_RAW),
RandomRotate90(0.2),
Flip(p=0.2),
ShiftScaleRotate(),
ToTensor()
])
from albumentations import Compose, RandomBrightnessContrast, \
HorizontalFlip, FancyPCA, HueSaturationValue, OneOf, ToGray, \
ShiftScaleRotate, ImageCompression, PadIfNeeded, GaussNoise, GaussianBlur
import cv2
from albu import IsotropicResize
size = 256
# Declare an augmentation pipeline
transform = Compose([
ImageCompression(quality_lower=60, quality_upper=100, p=0.5),
GaussNoise(p=0.1),
GaussianBlur(blur_limit=3, p=0.05),
HorizontalFlip(p=0.5),
OneOf([
IsotropicResize(max_side=size, interpolation_down=cv2.INTER_AREA, interpolation_up=cv2.INTER_CUBIC),
IsotropicResize(max_side=size, interpolation_down=cv2.INTER_AREA, interpolation_up=cv2.INTER_LINEAR),
IsotropicResize(max_side=size, interpolation_down=cv2.INTER_LINEAR, interpolation_up=cv2.INTER_LINEAR),
], p=0.7),
PadIfNeeded(min_height=size, min_width=size, border_mode=cv2.BORDER_CONSTANT),
OneOf([RandomBrightnessContrast(), FancyPCA(), HueSaturationValue()], p=0.7),
ToGray(p=0.2),
ShiftScaleRotate(shift_limit=0.1, scale_limit=0.2, rotate_limit=10, border_mode=cv2.BORDER_CONSTANT, p=0.5),
])
# Read an image with OpenCV and convert it to the RGB colorspace
image = cv2.imread("/home/ubuntu/dataset/dfdc_image/train/dfdc_train_part_0/aaqaifqrwn/frame0.jpg")
#image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
# Augment an image
transformed = transform(image=image)
def __init__(self,
normalization: Optional[str] = None,
whitespace_normalization: bool = True,
reorder: bool = True,
im_transforms: Callable[[Any],
torch.Tensor] = transforms.Compose(
[]),
preload: bool = True,
augmentation: bool = False) -> None:
self._images = [] # type: Union[List[Image], List[torch.Tensor]]
self._gt = [] # type: List[str]
self.alphabet = Counter() # type: Counter
self.text_transforms = [] # type: List[Callable[[str], str]]
# split image transforms into two. one part giving the final PIL image
# before conversion to a tensor and the actual tensor conversion part.
self.head_transforms = transforms.Compose(im_transforms.transforms[:2])
self.tail_transforms = transforms.Compose(im_transforms.transforms[2:])
self.transforms = im_transforms
self.preload = preload
self.aug = None
self.seg_type = 'baselines'
# built text transformations
if normalization:
self.text_transforms.append(
lambda x: unicodedata.normalize(cast(str, normalization), x))
if whitespace_normalization:
self.text_transforms.append(
lambda x: regex.sub('\s', ' ', x).strip())
if reorder:
self.text_transforms.append(bd.get_display)
if augmentation:
from albumentations import (
Compose,
ToFloat,
FromFloat,
Flip,
OneOf,
MotionBlur,
MedianBlur,
Blur,
ShiftScaleRotate,
OpticalDistortion,
ElasticTransform,
RandomBrightnessContrast,
)
self.aug = Compose([
ToFloat(),
OneOf([
MotionBlur(p=0.2),
MedianBlur(blur_limit=3, p=0.1),
Blur(blur_limit=3, p=0.1),
],
p=0.2),
ShiftScaleRotate(
shift_limit=0.0625, scale_limit=0.2, rotate_limit=3,
p=0.2),
OneOf([
OpticalDistortion(p=0.3),
ElasticTransform(p=0.1),
],
p=0.2),
],
p=0.5)
self.im_mode = '1'
def main(argv):
try:
opts, args = getopt.getopt(argv, "hi:o:", ["ifile=", "ofile="])
except getopt.GetoptError:
print('example.py -i <inputfile> -o <outputfile>')
sys.exit(2)
for opt, arg in opts:
if opt == '-h':
print('example.py -i <inputfile> -o <outputfile>')
sys.exit()
elif opt in ("-i", "--ifile"):
inputpath = arg
elif opt in ("-o", "--ofile"):
outputpath = arg
print('Input file is ', inputpath)
print('Output file is ', outputpath)
try:
os.mkdir(outputpath)
except OSError:
print("Folder %s already exists " % outputpath)
else:
print("Successfully created the directory %s " % outputpath)
for dirnames in os.listdir(inputpath):
print("Current name is: ", inputpath + dirnames)
if (os.path.isdir(inputpath + dirnames)):
dirname_new = dirnames
try:
os.mkdir(outputpath + dirname_new)
except OSError:
print("Folder %s already exists " % outputpath + dirname_new)
else:
print("Successfully created the directory %s " % outputpath +
dirname_new)
for filename in os.listdir(inputpath + dirnames):
whole_path = inputpath + dirnames + '/' + filename
if (whole_path.find('.png') != -1):
image = load_rgb_image(inputpath + dirnames + '/' + filename)
# aug = HorizontalFlip(p=1)
# aug = IAAAdditiveGaussianNoise(p=1)
# aug = RandomBrightnessContrast(p=1)
# aug = Blur(blur_limit=(5, 5),p=1)
# aug = CLAHE(p=1)
# aug = IAASharpen(p=1)
aug = Compose([
OneOf([
IAAAdditiveGaussianNoise(p=0.5),
RandomBrightnessContrast(p=0.5)
], 0.75),
OneOf([Blur(blur_limit=(5, 5), p=0.5),
CLAHE(p=0.5)], 0.75)
])
image = aug(image=image)['image']
# aug2 = CLAHE(p=1)
# image = aug2(image=image)['image']
cv2.imwrite(outputpath + dirname_new + '/' + filename,
cv2.cvtColor(image, cv2.COLOR_RGB2BGR))
else:
print('wrong file type')
def __init__(self,
split: Callable[[str], str] = lambda x: path.splitext(x)[0],
suffix: str = '.gt.txt',
normalization: Optional[str] = None,
whitespace_normalization: bool = True,
reorder: bool = True,
im_transforms: Callable[[Any],
torch.Tensor] = transforms.Compose(
[]),
preload: bool = True,
augmentation: bool = False) -> None:
"""
Reads a list of image-text pairs and creates a ground truth set.
Args:
split (func): Function for generating the base name without
extensions from paths
suffix (str): Suffix to attach to image base name for text
retrieval
mode (str): Image color space. Either RGB (color) or L
(grayscale/bw). Only L is compatible with vertical
scaling/dewarping.
scale (int, tuple): Target height or (width, height) of dewarped
line images. Vertical-only scaling is through
CenterLineNormalizer, resizing with Lanczos
interpolation. Set to 0 to disable.
normalization (str): Unicode normalization for gt
whitespace_normalization (str): Normalizes unicode whitespace and
strips whitespace.
reorder (bool): Whether to rearrange code points in "display"/LTR
order
im_transforms (func): Function taking an PIL.Image and returning a
tensor suitable for forward passes.
preload (bool): Enables preloading and preprocessing of image files.
"""
self.suffix = suffix
self.split = lambda x: split(x) + self.suffix
self._images = [] # type: Union[List[Image], List[torch.Tensor]]
self._gt = [] # type: List[str]
self.alphabet = Counter() # type: Counter
self.text_transforms = [] # type: List[Callable[[str], str]]
# split image transforms into two. one part giving the final PIL image
# before conversion to a tensor and the actual tensor conversion part.
self.head_transforms = transforms.Compose(im_transforms.transforms[:2])
self.tail_transforms = transforms.Compose(im_transforms.transforms[2:])
self.aug = None
self.preload = preload
self.seg_type = 'bbox'
# built text transformations
if normalization:
self.text_transforms.append(
lambda x: unicodedata.normalize(cast(str, normalization), x))
if whitespace_normalization:
self.text_transforms.append(
lambda x: regex.sub('\s', ' ', x).strip())
if reorder:
self.text_transforms.append(bd.get_display)
if augmentation:
from albumentations import (
Compose,
ToFloat,
FromFloat,
Flip,
OneOf,
MotionBlur,
MedianBlur,
Blur,
ShiftScaleRotate,
OpticalDistortion,
ElasticTransform,
RandomBrightnessContrast,
)
self.aug = Compose([
ToFloat(),
OneOf([
MotionBlur(p=0.2),
MedianBlur(blur_limit=3, p=0.1),
Blur(blur_limit=3, p=0.1),
],
p=0.2),
ShiftScaleRotate(shift_limit=0.0625,
scale_limit=0.2,
rotate_limit=45,
p=0.2),
OneOf([
OpticalDistortion(p=0.3),
ElasticTransform(p=0.1),
],
p=0.2),
],
p=0.5)
self.im_mode = '1'
def get_transforms(phase_config):
list_transforms = []
if phase_config.Resize.p > 0:
list_transforms.append(
Resize(phase_config.Resize.height, phase_config.Resize.width, p=1))
if phase_config.HorizontalFlip:
list_transforms.append(HorizontalFlip())
if phase_config.VerticalFlip:
list_transforms.append(VerticalFlip())
if phase_config.RandomCropScale:
if phase_config.Resize.p > 0:
height = phase_config.Resize.height
width = phase_config.Resize.width
else:
height = HEIGHT
width = WIDTH
list_transforms.append(
RandomSizedCrop(min_max_height=(int(height * 0.90), height),
height=height,
width=width,
w2h_ratio=width / height))
if phase_config.ShiftScaleRotate:
list_transforms.append(ShiftScaleRotate(p=1))
if phase_config.RandomCrop.p > 0:
list_transforms.append(
RandomCrop(phase_config.RandomCrop.height,
phase_config.RandomCrop.width,
p=1))
if phase_config.Noise:
list_transforms.append(
OneOf([
GaussNoise(),
IAAAdditiveGaussianNoise(),
], p=0.5), )
if phase_config.Contrast:
list_transforms.append(
OneOf([
RandomContrast(0.5),
RandomGamma(),
RandomBrightness(),
],
p=0.5), )
if phase_config.Blur:
list_transforms.append(
OneOf([
MotionBlur(p=.2),
MedianBlur(blur_limit=3, p=0.1),
Blur(blur_limit=3, p=0.1),
],
p=0.5))
if phase_config.Distort:
list_transforms.append(
OneOf([
OpticalDistortion(p=0.3),
GridDistortion(p=.1),
IAAPiecewiseAffine(p=0.3),
],
p=0.5))
if phase_config.Cutout.num_holes > 0:
num_holes = phase_config.Cutout.num_holes
hole_size = phase_config.Cutout.hole_size
list_transforms.append(Cutout(num_holes, hole_size))
list_transforms.extend([
Normalize(mean=phase_config.mean, std=phase_config.std, p=1),
ToTensor(),
])
return Compose(list_transforms)
RandomBrightnessContrast, RandomContrast, RandomGamma, OneOf, ToFloat,
ShiftScaleRotate, GridDistortion, ElasticTransform, JpegCompression,
HueSaturationValue, RGBShift, RandomBrightness, RandomContrast, Blur,
MotionBlur, MedianBlur, GaussNoise, CenterCrop, IAAAdditiveGaussianNoise,
GaussNoise, OpticalDistortion)
AUGMENTATIONS_TRAIN = Compose(
[
HorizontalFlip(p=0.5),
#optionas below emulate different positioning and patient geometry
ShiftScaleRotate(rotate_limit=(-25, 25)),
#options below emulate different X-ray penetration
OneOf([
RandomContrast(),
RandomGamma(),
RandomBrightness(),
Blur(blur_limit=2),
],
p=0.3),
#ToFloat(max_value=1)
],
p=1)
AUGMENTATIONS_TEST = Compose(
[
#ToFloat(max_value=1)
],
p=1)
class generator(keras.utils.Sequence):
def __init__(self, opt):
# super(Batch_Balanced_Dataset, self).__init__()
self.tbaug = TenBrightAug()
self.incbaug = IncBrightAug()
self.colaug = ColorAug()
self.distortaug = DistortAug()
self.grayimg = GrayImg()
self.binimg = BinImg()
self.jpgaug = JPEGAug(0.8)
self.resizeimg = ResizeAug(opt.imgW, opt.imgH)
self.resizeimg_test = ResizeAug(opt.imgW, opt.imgH, rand_scale=False)
l1 = len(opt.train_data.strip().split(','))
if l1 == 1:
self.batch_sizes = [int(opt.batch_size)]
elif l1 == 3:
self.batch_sizes = [
int(opt.batch_size * opt.batch_ratio),
opt.batch_size - int(opt.batch_size * opt.batch_ratio)
]
elif l1 == 4:
b1 = int(opt.batch_size * opt.batch_ratio2)
self.batch_sizes = [int(b1 * opt.batch_ratio), 0, 0]
self.batch_sizes[1] = b1 - self.batch_sizes[0]
self.batch_sizes[
2] = opt.batch_size - self.batch_sizes[0] - self.batch_sizes[1]
if not opt.alldata:
self.test_books = [
'200021660', '200005598', 'hnsd006', 'umgy003', '100249416',
'umgy011', 'umgy010'
]
else:
print('use all data')
self.test_books = [
# '200021660',
# '200005598',
'hnsd006',
'umgy003',
'100249416',
'umgy011',
'umgy010'
]
self.train_tf1 = transforms.Compose([
self.distortaug, self.colaug, self.tbaug, self.incbaug,
self.grayimg, self.binimg, self.tbaug, self.incbaug, self.jpgaug,
self.resizeimg,
transforms.ToTensor()
])
self.train_tf0 = transforms.Compose(
[self.grayimg, self.resizeimg,
transforms.ToTensor()])
self.train_tf00 = Compose([
OneOf([
OpticalDistortion(distort_limit=0.05,
shift_limit=10,
border_mode=cv2.BORDER_WRAP,
p=0.5),
ElasticTransform(p=0.5,
alpha=1,
sigma=50,
alpha_affine=0.2,
border_mode=cv2.BORDER_CONSTANT)
]),
OneOf([
CLAHE(),
Solarize(),
RandomBrightness(),
RandomContrast(limit=0.2),
RandomBrightnessContrast(),
]),
JpegCompression(quality_lower=20, quality_upper=100),
RandomShadow(num_shadows_lower=1, num_shadows_upper=2),
PadIfNeeded(min_height=64,
min_width=100,
border_mode=cv2.BORDER_CONSTANT,
p=0.5),
Cutout(num_holes=8, max_h_size=16, max_w_size=16),
InvertImg(p=0.3),
ToGray()
])
self.train_tf01 = Compose([
JpegCompression(quality_lower=20, quality_upper=100),
Cutout(num_holes=4, max_h_size=8, max_w_size=8),
InvertImg(p=0.3)
])
self.randtf = transforms.RandomApply([self.distortaug])
self.train_tf2 = transforms.Compose(
[self.resizeimg_test, transforms.ToTensor()])
self.data_loader_list = []
self.datasets = []
self.data_samplers = []
self.opt = opt
self.train_data = opt.train_data.strip().split(',')
if not 'txt' in self.train_data[0]:
self.datasets.append(
RealImageLMDB(self.train_data[0],
transform=self.train_tf0,
transform2=self.train_tf00,
testBooks=self.test_books,
character=opt.character,
max_batch_length=opt.batch_max_length))
else:
self.datasets.append(
RealImageTxtLMDB(self.train_data[0],
transform=self.train_tf0,
transform2=self.train_tf00,
testBooks=self.test_books,
max_batch_length=opt.batch_max_length))
if len(self.train_data) == 3:
self.datasets.append(
SynthImageLMDB(self.train_data[1],
self.train_data[2],
transform=self.train_tf2,
transform2=self.train_tf01,
size=(opt.imgH, opt.imgW),
gray_bin_ratio=0.5,
max_label_length=opt.batch_max_length))
if len(self.train_data) == 4:
self.datasets.append(
RealImageTxtLMDB(self.train_data[3],
transform=self.train_tf0,
transform2=self.train_tf00,
testBooks=self.test_books,
max_batch_length=opt.batch_max_length))
for i in range(len(self.datasets)):
train_sampler = torch.utils.data.distributed.DistributedSampler(
self.datasets[i])
self.data_samplers.append(train_sampler)
self.data_loader_list.append(
DataLoader(self.datasets[i],
num_workers=int(opt.workers),
shuffle=False,
sampler=train_sampler,
batch_size=self.batch_sizes[i],
collate_fn=collate_fn,
pin_memory=True,
drop_last=True))
self.dataloader_iter_list = [iter(i) for i in self.data_loader_list]
self.test_tf = transforms.Compose(
[self.grayimg, self.resizeimg_test,
transforms.ToTensor()])
self.test_dataset = RealImageLMDB(self.train_data[0],
self.test_tf,
testBooks=self.test_books,
isTest=True,
character=opt.character)
self.test_sampler = torch.utils.data.distributed.DistributedSampler(
self.test_dataset)
self.test_loader = DataLoader(self.test_dataset,
num_workers=int(opt.workers),
shuffle=False,
sampler=self.test_sampler,
batch_size=max(2,
int(opt.batch_size / 8)),
collate_fn=collate_fn,
drop_last=True)
import numpy as np
from albumentations import (
Compose, HorizontalFlip, CLAHE, HueSaturationValue,
RandomBrightness, RandomContrast, RandomGamma,OneOf,
ToFloat, ShiftScaleRotate,GridDistortion, ElasticTransform, JpegCompression, HueSaturationValue,
RGBShift, RandomBrightness, RandomContrast, Blur, MotionBlur, MedianBlur, GaussNoise,CenterCrop,
IAAAdditiveGaussianNoise,GaussNoise,OpticalDistortion,RandomSizedCrop
)
size = 512
AUGMENTATIONS_TRAIN = Compose([
HorizontalFlip(p=0.5),
OneOf([
RandomContrast(),
RandomGamma(),
RandomBrightness()
], p=0.3),
OneOf([
ElasticTransform(alpha = 120, sigma=120*0.05,alpha_affine = 12*0.03),
GridDistortion(),
OpticalDistortion(distort_limit = 2, shift_limit = 0.5)
], p=0.3),
RandomSizedCrop(min_max_height=(512,1024),height = size, width =size,p=1),
ToFloat(max_value=1)
], p =1)
AUGMENTATIONS_TEST = Compose([
RandomSizedCrop(min_max_height=(512,1024),height = size, width =size,p=1),
ToFloat(max_value=1)
],p=1)
AUGMENTATIONS_TEST2 = Compose([
ToFloat(max_value=1)
if __name__ == '__main__':
import matplotlib.pyplot as plt
from tqdm import tqdm
from albumentations import Compose, OneOf, Normalize
from albumentations import HorizontalFlip, VerticalFlip, RandomRotate90, RandomScale, RandomCrop
from albumentations.pytorch import ToTensorV2
dataset = Xview2(
r'D:\DATA\xView2\train\images',
r'D:\DATA\xView2\train\labels',
transforms=Compose([
OneOf([
HorizontalFlip(True),
VerticalFlip(True),
RandomRotate90(True)
],
p=0.5),
# RandomDiscreteScale([0.75, 1.25, 1.5], p=0.5),
RandomCrop(640, 640, True),
Normalize(mean=(0.485, 0.456, 0.406, 0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225, 0.229, 0.224, 0.225),
max_pixel_value=255),
ToTensorV2(True),
]),
include=('pre', 'post')).pairwise_mode()
print(len(dataset))
a = dataset[1]
print()
# img, mask = dataset[4]
def forward(self, image):
out = self.base_model(image)
return out
from albumentations import (
HorizontalFlip, IAAPerspective, ShiftScaleRotate, CLAHE, RandomRotate90,
Transpose, ShiftScaleRotate, Blur, OpticalDistortion, GridDistortion,
HueSaturationValue, IAAAdditiveGaussianNoise, GaussNoise, MotionBlur,
MedianBlur, RandomBrightnessContrast, IAAPiecewiseAffine, IAASharpen,
IAAEmboss, Flip, OneOf, Compose, Rotate, Cutout, HorizontalFlip, Normalize)
from albumentations.pytorch import ToTensor
train_aug = Compose([
Rotate(15),
OneOf([
IAAAdditiveGaussianNoise(),
GaussNoise(),
], p=0.2),
OneOf([
MotionBlur(p=0.2),
MedianBlur(blur_limit=3, p=0.1),
Blur(blur_limit=3, p=0.1),
],
p=0.2),
ShiftScaleRotate(shift_limit=0.0625,
scale_limit=0.2,
rotate_limit=45,
p=0.2),
OneOf([
OpticalDistortion(p=0.3),
GridDistortion(p=0.1),
IAAPiecewiseAffine(p=0.3),
def initialize(self, config, filename):
self.config = config
self.train = 'train' in filename
self.size = config['image_size']
if self.train:
self.transform = Compose([
Resize(self.size[0], self.size[1]),
ShiftScaleRotate(shift_limit=0.3,
scale_limit=(0.05, 0.1),
rotate_limit=10,
p=.4),
OneOf([
IAAAdditiveGaussianNoise(),
GaussNoise(),
], p=0.4),
OneOf([
MotionBlur(p=.2),
MedianBlur(blur_limit=5, p=.5),
Blur(blur_limit=3, p=.5),
],
p=0.4),
OpticalDistortion(p=0.4),
OneOf([
CLAHE(clip_limit=3),
IAASharpen(),
IAAEmboss(),
RandomContrast(),
RandomBrightness(),
RandomGamma()
],
p=0.6),
OneOf([
RGBShift(),
HueSaturationValue(),
], p=0.2),
JpegCompression(quality_lower=30, quality_upper=100, p=0.4),
Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225],
),
])
else:
self.transform = Compose([
Resize(self.size[0], self.size[1]),
Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225],
),
])
with open(filename, 'r') as f:
reader = csv.reader(f)
self.anno = list(reader)
with open(self.config['datasets']['bboxs'], 'r') as bboxs_file:
bboxs_json = json.load(bboxs_file)
self.bboxs_dict = {}
for item in bboxs_json:
if item['image_id'] in self.bboxs_dict:
fixed_bbox_new, area_new = self.fix_bbox(item['bbox'])
_, area_old = self.fix_bbox(
self.bboxs_dict[item['image_id']])
if area_new > area_old:
self.bboxs_dict[item['image_id']] = fixed_bbox_new
else:
self.bboxs_dict[item['image_id']] = self.fix_bbox(
item['bbox'])[0]
random.shuffle(self.anno)
self.anno_size = len(self.anno)
def load_data(layer, val = False):
assert(isinstance(layer, str))
TRAIN = 512
TRAIN_ADD = 384
VALID_IN = 50
VALID_X = 39
PATH = '/content/drive/My Drive/Rosneft Seismic/train/'
augm = OneOf([
Cutout(num_holes=4, max_h_size=38, max_w_size=38, always_apply=True),
Cutout(num_holes=50, max_h_size=6, max_w_size=6, always_apply=True),
Cutout(num_holes=20, max_h_size=15, max_w_size=15, always_apply=True)
], p=1.)
x = []
y = []
borders = []
if val:
inds_in = list(itertools.chain(range(0, VALID_IN), range(TRAIN-VALID_IN, TRAIN)))
inds_x = list(itertools.chain(range(0, VALID_X), range(TRAIN_ADD-VALID_X, TRAIN_ADD)))
print('LOAD VALIDATION DATA')
else:
inds_in = list(range(VALID_IN, TRAIN-VALID_IN))
inds_x = list(range(VALID_X, TRAIN_ADD-VALID_X))
print('LOAD TRAIN DATA')
print('--load inline images')
for i in tqdm_notebook(inds_in):
img = np.array(Image.open(PATH+'images/inline_'+str(i+3160)+'.png'))
mask = np.array(Image.open(PATH+'answer/'+layer+'/inline_'+str(i+3160)+'.png'))
border = np.array(Image.open(PATH+'answer/'+layer+'_border/inline_'+str(i+3160)+'.png'))
img = img.astype('float32')
img = img[:, :, 0]*0.299 + img[:, :, 1]*0.587 + img[:, :, 2]*0.114
img -= img.mean()
img /= img.std()
x.append(img)
y.append(mask)
borders.append(border)
if not val:
x.append(img[:, ::-1])
y.append(mask[:, ::-1])
borders.append(border[:, ::-1])
x.append(augm(image=img)["image"])
y.append(mask)
borders.append(border)
print('--load xline images')
for i in tqdm_notebook(inds_x):
img = np.array(Image.open(PATH+'images/xline_'+str(i+2017)+'.png'))
mask = np.array(Image.open(PATH+'answer/'+layer+'/xline_'+str(i+2017)+'.png'))
border = np.array(Image.open(PATH+'answer/'+layer+'_border/xline_'+str(i+2017)+'.png'))
img = img.astype('float32')
img = img[:, :, 0]*0.299 + img[:, :, 1]*0.587 + img[:, :, 2]*0.114
img_left = img[:, 0:384]
img_left -= img_left.mean()
img_left /= img_left.std()
img_right = img[:, 128:512]
img_right -= img_right.mean()
img_right /= img_right.std()
mask_left = mask[:, 0:384]
mask_right = mask[:, 128:512]
border_left = border[:, 0:384]
border_right = border[:, 128:512]
x.append(img_left)
x.append(img_right)
y.append(mask_left)
y.append(mask_right)
borders.append(border_left)
borders.append(border_right)
if not val:
x.append(img_left[:, ::-1])
x.append(img_right[:, ::-1])
y.append(mask_left[:, ::-1])
y.append(mask_right[:, ::-1])
borders.append(border_left[:, ::-1])
borders.append(border_right[:, ::-1])
x.append(augm(image=img_left)["image"])
x.append(augm(image=img_right)["image"])
y.append(mask_left)
y.append(mask_right)
borders.append(border_left)
borders.append(border_right)
if val:
pass
print('END LOAD VALIDATION DATA')
else:
print('END LOAD TRAIN DATA')
print('STAR PREPROCCECING')
x = np.array(x)
x = x[:, :, :, np.newaxis]
y = np.array(y)
borders = np.array(borders)
y = y[:, :, :, np.newaxis, np.newaxis]
borders = borders[:, :, :, np.newaxis, np.newaxis]
y = np.concatenate((y, borders), axis = 3)
y = y.astype('float32')
y /= 255
print('END PREPROCCECING')
return x, y
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--train-image-dirs', required=True, nargs='+')
parser.add_argument('--train-mask-dirs', required=True, nargs='+')
parser.add_argument('--train-dataset-types', nargs='+')
parser.add_argument('--validation-image-dirs', required=True, nargs='+')
parser.add_argument('--validation-mask-dirs', required=True, nargs='+')
parser.add_argument('--validation-dataset-types', nargs='+')
parser.add_argument('--save-dir', default='./runs')
parser.add_argument('--cpu', action='store_true')
parser.add_argument('--category-type', default='binary', choices=['binary', 'simple'])
parser.add_argument('--epochs', type=int, default=10)
parser.add_argument('--input-size', type=int, nargs=2, default=(640, 640))
parser.add_argument('--jaccard-weight', type=float, default=0.3)
available_networks = ['unet11', 'unet16']
parser.add_argument('--model-name', type=str, default='unet11', choices=available_networks)
parser.add_argument('--batch-size', type=int, default=128)
parser.add_argument('--seed', type=int, default=1)
parser.add_argument('--device-id', type=int, default=0)
parser.add_argument('--run-name', type=str, required=True)
args = parser.parse_args()
print(args)
device = torch_tools.get_device(args.cpu, args.device_id)
torch_tools.set_seeds(args.seed, device)
train_image_dirs = [Path(p) for p in args.train_image_dirs]
train_mask_dirs = [Path(p) for p in args.train_mask_dirs]
validation_image_dirs = [Path(p) for p in args.validation_image_dirs]
validation_mask_dirs = [Path(p) for p in args.validation_mask_dirs]
for data_dir in\
train_image_dirs +\
train_mask_dirs +\
validation_image_dirs +\
validation_mask_dirs:
assert data_dir.exists(), f'{str(data_dir)} does not exist.'
input_size = args.input_size
w, h = input_size
rate = 0.9
'''
train_transform = Compose([
HorizontalFlip(p=0.5),
RandomCrop(*input_size),
])
'''
# Transforms
transforms = {}
# Basic transform
transforms['base'] = Compose([
HorizontalFlip(p=0.5),
#IAAPerspective(scale=(0.05, 0.1), p=0.3),
Rotate(5, p=0.5),
RandomGamma(p=0.5),
HueSaturationValue(
hue_shift_limit=10,
sat_shift_limit=15,
val_shift_limit=10,
p=0.5
),
RandomBrightnessContrast(p=0.5),
OneOf([
RandomSizedCrop((int(h * rate), int(w * rate)), h, w, p=1.0),
RandomCrop(h, w, p=1.0),
], p=1.0)
])
# BDD dataset
transforms['bdd'] = transforms['base']
# Always shrink to 22% - 50%
transforms['walk'] = Compose([
HorizontalFlip(p=0.5),
Rotate(5, p=0.5),
RandomGamma(p=0.5),
HueSaturationValue(
hue_shift_limit=10,
sat_shift_limit=15,
val_shift_limit=10,
p=0.5
),
RandomBrightnessContrast(p=0.5),
RandomScale((-0.78, -0.5), p=1.0),
RandomCrop(h, w, p=1.0),
])
# MISC dataset transform
transforms['misc'] = Compose([
HorizontalFlip(p=0.5),
Rotate(5, p=0.5),
RandomGamma(p=0.5),
HueSaturationValue(
hue_shift_limit=10,
sat_shift_limit=15,
val_shift_limit=10,
p=0.5
),
RandomBrightnessContrast(p=0.5),
Resize(h, w, p=1.0),
])
validation_transform = Compose([
CenterCrop(h, w),
])
category_type = datasets.surface_types.from_string(args.category_type)
# Logger
log_dir = _get_log_dir(args)
logger = logging.Logger(log_dir, n_save=16, image_size=256, category_type=category_type)
logger.writer.add_text('args', str(args))
train_datasets = []
for image_dir, mask_dir, dataset_type in zip(train_image_dirs, train_mask_dirs, args.train_dataset_types):
_dataset = datasets.create_dataset(
dataset_type,
[image_dir],
[mask_dir],
category_type,
transforms[dataset_type],
)
train_datasets.append(_dataset)
validation_datasets = []
for image_dir, mask_dir, dataset_type in zip(validation_image_dirs, validation_mask_dirs, args.validation_dataset_types):
_dataset = datasets.create_dataset(
dataset_type,
[image_dir],
[mask_dir],
category_type,
validation_transform,
)
validation_datasets.append(_dataset)
# Merge datasets
train_dataset = ConcatDataset(train_datasets)
validation_dataset = ConcatDataset(validation_datasets)
train_loader = DataLoader(train_dataset, args.batch_size, shuffle=True)
validation_loader = DataLoader(validation_dataset, args.batch_size, shuffle=False)
net = models.load_model(args.model_name, category_type).to(device)
criterion = models.loss.get_criterion(category_type, args.jaccard_weight)
optimizer = torch.optim.Adam(net.parameters())
for epoch in range(1, args.epochs + 1):
print(f'epoch: {epoch:03d}')
sys.stdout.flush()
train(net, train_loader, epoch, optimizer, criterion, device, logger)
evaluate(net, validation_loader, epoch, criterion, device, logger, 'validation')
temp = 0
for im in images_train:
if 'temp' in im:
temp += 1
cw = np.array([temp, len(images_train) - temp])
cw = np.median(cw / sum(cw)) / (cw / sum(cw))
class_weights = {0: cw[0], 1: cw[1]}
# In[]: Augmentor
aug = OneOf([
Blur(blur_limit=5, p=1.),
RandomGamma(gamma_limit=(50, 150), p=1.),
HueSaturationValue(
hue_shift_limit=20, sat_shift_limit=30, val_shift_limit=20, p=1.),
RGBShift(r_shift_limit=15, g_shift_limit=5, b_shift_limit=15, p=1.),
RandomBrightness(limit=.25, p=1.),
RandomContrast(limit=.25, p=1.),
MedianBlur(blur_limit=5, p=1.),
CLAHE(clip_limit=2.0, tile_grid_size=(8, 8), p=1.)
],
p=.5)
def augment(image, aug=aug):
augmented = aug(image=image)
image_augmented = augmented['image']
return image_augmented
# In[]: Custom generator
def custom_generator(images_path, batch_size=1, validate=False):
def train(
train_df,
valid_df,
model,
multi_model,
input_size,
epochs,
batch_size,
save_weights_path,
save_logs_path,
train_dir,
valid_dir,
):
if not os.path.exists(save_weights_path):
os.makedirs(save_weights_path)
if not os.path.exists(save_logs_path):
os.makedirs(save_logs_path)
AUGMENTATIONS_TRAIN = Compose([
HorizontalFlip(p=0.35),
RandomSizedCrop(min_max_height=(int(input_size * 0.75), input_size),
height=input_size,
width=input_size,
p=0.35),
OneOf([
ShiftScaleRotate(rotate_limit=25, p=0.35),
ElasticTransform(
alpha=120, sigma=120 * 0.05, alpha_affine=120 * 0.03, p=0.35),
GridDistortion(p=0.35),
OpticalDistortion(distort_limit=2, shift_limit=0.5, p=0.35),
],
p=1),
OneOf([
RandomContrast(p=0.35),
RandomGamma(p=0.35),
RandomBrightness(p=0.35)
],
p=1),
OneOf([
Blur(p=0.35),
MedianBlur(p=0.35),
GaussNoise(p=0.35),
GaussianBlur(p=0.35)
],
p=1)
],
p=1)
#Generator Parmas
params_train = {
'list_IDs': list(train_df.Image),
'labels': list(train_df.Labels),
'dim': (input_size, input_size),
'data_dir': train_dir,
'batch_size': batch_size,
'n_channels': 3,
'n_classees': PARAMS_JSON['CLASS_NUM'],
'aug': AUGMENTATIONS_TRAIN,
'model_name': None,
'preprocess_input': preprocess_input,
'to_categori': to_categori,
'shuffle': True
}
params_val = {
'list_IDs': list(valid_df.Image),
'labels': list(valid_df.Labels),
'dim': (input_size, input_size),
'data_dir': valid_dir,
'batch_size': batch_size,
'n_channels': 3,
'n_classees': PARAMS_JSON['CLASS_NUM'],
'aug': None,
'model_name': None,
'preprocess_input': preprocess_input,
'to_categori': to_categori,
'shuffle': True
}
#Create Generator
train_generator = DataGenerator(**params_train)
validation_generator = DataGenerator(**params_val)
#get class weight
class_weight_dict = CalculateClassWeight(train_df,
PARAMS_JSON['CLASS_NUM'],
to_categori)
#model check point
model_path = os.path.join(save_weights_path, '{}.h5')
# check_point = ModelCheckpoint(filepath=model_path, monitor='val_auc',verbose=1,mode=max)
#lr schedule
reduceLR = ReduceLROnPlateau(monitor='val_loss',
factor=0.5,
mode=min,
patience=5)
#callbacks of CsvLogger
logs_path = os.path.join(save_logs_path, "log.csv")
csvlogger = CSVLogger(logs_path)
cbk = MyCbk(model, model_path)
callbacks = [csvlogger, reduceLR, cbk]
#Train
multi_model.fit_generator(generator=train_generator,
epochs=epochs,
validation_data=validation_generator,
callbacks=callbacks,
class_weight=class_weight_dict)
from albumentations import Compose, Resize, RandomCrop, Flip, HorizontalFlip, VerticalFlip, Transpose, RandomRotate90, \
ShiftScaleRotate, OneOf, OpticalDistortion, HueSaturationValue, RandomGamma
from albumentations.pytorch import ToTensor
train_aug = Compose([
RandomCrop(height=96, width=96, p=0.2),
OneOf([
VerticalFlip(p=0.2),
HorizontalFlip(p=0.3),
Transpose(p=0.2),
RandomRotate90(p=0.2),
],
p=0.3),
ShiftScaleRotate(p=0.2),
RandomGamma(p=0.2),
Resize(128, 128, always_apply=True),
ToTensor()
])
valid_aug = Compose([Resize(128, 128, always_apply=True), ToTensor()])
def __init__(self, config):
super(AugmentedPair2, self).__init__(config)
self.use_appearance_augmentation = config.get("data_augment_appearance", False)
self.use_shape_augmentation = config.get("data_augment_shape", False)
additional_targets = {
"image{}".format(i): "image" for i in range(1, self.n_images)
}
p = 0.9
appearance_augmentation = Compose(
[
OneOf(
[
MedianBlur(blur_limit=3, p=0.1),
Blur(blur_limit=3, p=0.1),
],
p=0.5,
),
OneOf(
[
RandomBrightnessContrast(p=0.3),
RGBShift(p=0.3),
HueSaturationValue(p=0.3),
],
p=0.8,
),
OneOf(
[
RandomBrightnessContrast(p=0.3),
RGBShift(p=0.3),
HueSaturationValue(p=0.3),
],
p=0.8,
),
OneOf(
[
RandomBrightnessContrast(p=0.3),
RGBShift(p=0.3),
HueSaturationValue(p=0.3),
],
p=0.8,
),
ToGray(p=0.1),
ChannelShuffle(p=0.3),
],
p=p,
additional_targets=additional_targets,
)
self.appearance_augmentation = appearance_augmentation
p = 0.9
shape_augmentation = Compose(
[
HorizontalFlip(p=0.3),
ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.25, rotate_limit=25, p=0.3, border_mode=cv2.BORDER_REPLICATE),
OneOf([
# OpticalDistortion(p=0.3),
# GridDistortion(p=0.1),
IAAPiecewiseAffine(p=0.5),
ElasticTransform(p=0.5, border_mode=cv2.BORDER_REPLICATE)
], p=0.3),
],
p=p,
additional_targets=additional_targets,
)
self.shape_augmentation = shape_augmentation