需要注意增强之后以字典键值的形式调用输出,且可以保证img、gt的随机增强是一一对应一致的。
import albumentations as albu
tfrd = albu.Compose([
albu.HorizontalFlip(p=0.5),
albu.VerticalFlip(p=0.5),
albu.ShiftScaleRotate(shift_limit = 0,scale_limit = 0,rotate_limit=(-30,30),p=1,border_mode=cv2.BORDER_CONSTANT,value=0),
albu.Blur(blur_limit=3, p=0.3),
])
for name in trainlist:
img = cv2.imread(os.path.join(img_dir,name))
gt = cv2.imread(os.path.join(gt_dir,name))
for j in range(4):
ns = tfrd(image=img,mask=gt)
newimg,newgt= ns['image'],ns['mask']
cv2.imwrite(os.path.join(img_dir,'t'+str(j)+'_'+name),newimg)
cv2.imwrite(os.path.join(gt_dir,'t'+str(j)+'_'+name),newgt)
tmp.append('t'+str(j)+'_'+name)