def __init__(self):
transformation = []
transformation += [
RandomResizedCrop(512, 512, scale=(0.8, 1.0)),
ShiftScaleRotate(),
OneOf([
GridDistortion(),
OpticalDistortion(),
ElasticTransform(approximate=True)
],
p=0.8),
OneOf([
IAAAdditiveGaussianNoise(),
GaussNoise(),
MedianBlur(blur_limit=3),
Blur(blur_limit=3)
],
p=0.8),
CoarseDropout(max_holes=2, max_height=32, max_width=32)
]
self.transform = Compose(transformation)
def get_train_transforms():
return Compose([
RandomResizedCrop(config.H, config.W),
Transpose(p=0.5),
HorizontalFlip(p=0.5),
VerticalFlip(p=0.5),
ShiftScaleRotate(p=0.5),
HueSaturationValue(hue_shift_limit=0.2,
sat_shift_limit=0.2,
val_shift_limit=0.2,
p=0.5),
RandomBrightnessContrast(
brightness_limit=(-0.1, 0.1), contrast_limit=(-0.1, 0.1), p=0.5),
Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225],
max_pixel_value=255.0,
p=1.0),
CoarseDropout(p=0.5),
Cutout(p=0.5),
ToTensorV2(p=1.0),
],
p=1.)
def __init__(self, dataframe, vertical_flip, horizontal_flip,
is_train=True):
self.dataframe, self.is_train = dataframe, is_train
self.vertical_flip, self.horizontal_flip = vertical_flip, horizontal_flip
# Data Augmentation (custom for each dataset type)
if is_train:
self.transform = Compose([RandomResizedCrop(height=224, width=224, scale=(0.7, 1.0)),
ShiftScaleRotate(rotate_limit=90, scale_limit = [0.7, 1]),
HorizontalFlip(p = self.horizontal_flip),
VerticalFlip(p = self.vertical_flip),
HueSaturationValue(sat_shift_limit=[0.7, 1.3],
hue_shift_limit=[-0.1, 0.1]),
RandomBrightnessContrast(brightness_limit=[0.01, 0.1],
contrast_limit= [0.01, 0.1]),
#Normalize(),
ToTensor()])
else:
self.transform = Compose([ # Normalize(),
Resize(height=224, width=224),
ToTensor()])
def get_transforms(*, data):
if data == 'train':
return Compose([
# Resize(CFG.size, CFG.size),
RandomResizedCrop(CFG.size, CFG.size, scale=(0.85, 1.0)),
HorizontalFlip(p=0.5),
RandomBrightnessContrast(p=0.2,
brightness_limit=(-0.2, 0.2),
contrast_limit=(-0.2, 0.2)),
HueSaturationValue(p=0.2,
hue_shift_limit=0.2,
sat_shift_limit=0.2,
val_shift_limit=0.2),
ShiftScaleRotate(p=0.2,
shift_limit=0.0625,
scale_limit=0.2,
rotate_limit=20),
CoarseDropout(p=0.2),
Cutout(p=0.2,
max_h_size=16,
max_w_size=16,
fill_value=(0., 0., 0.),
num_holes=16),
Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225],
),
ToTensorV2(),
])
elif data == 'valid':
return Compose([
Resize(CFG.size, CFG.size),
Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225],
),
ToTensorV2(),
])
def build_aun_fn(image_size):
transforms = Compose([
RandomResizedCrop(image_size, image_size),
Transpose(p=0.5),
HorizontalFlip(p=0.5),
VerticalFlip(p=0.5),
ShiftScaleRotate(p=0.5),
HueSaturationValue(hue_shift_limit=0.2,
sat_shift_limit=0.2,
val_shift_limit=0.2,
p=0.5),
RandomBrightnessContrast(
brightness_limit=(-0.1, 0.1), contrast_limit=(-0.1, 0.1), p=0.5),
Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225],
max_pixel_value=255.0,
p=1.0),
CoarseDropout(p=0.5),
Cutout(p=0.5)
],
p=1.)
def aug_fn(image):
data = {"image": image}
aug_data = transforms(**data)
aug_img = aug_data["image"]
aug_img = tf.cast(aug_img / 255.0, tf.float32)
aug_img = tf.image.resize(aug_img, size=[image_size, image_size])
return aug_img
def augment(image, label):
aug_img = tf.numpy_function(func=aug_fn, inp=[image], Tout=tf.float32)
return aug_img, label
return augment
def train_augments(augments_config, mean, std, image_size):
box_scale = min(image_size)
return Compose([
HorizontalFlip(p=augments_config["HORIZONTAL_FLIP_PR"]),
VerticalFlip(p=augments_config["VERTICAL_FLIP_PR"]),
RandomResizedCrop(image_size[0],
image_size[1],
scale=(1, 1 + augments_config["RESCALE_PERCENTAGE"]),
ratio=(image_size[1] / image_size[0],
image_size[1] / image_size[0]),
interpolation=INTER_LINEAR,
p=augments_config["RESCALE_PR"]),
ElasticTransform(alpha=box_scale,
sigma=box_scale * 0.05,
alpha_affine=box_scale * 0.03,
p=augments_config["ELASTIC_TRANSFORM_PR"]),
RandomGamma(gamma_limit=(80, 120), p=augments_config["GAMMA_PR"]),
Normalize(max_pixel_value=1.0,
mean=mean,
std=std,
p=augments_config["NORMALIZE_PR"])
],
p=1)
def build_train(self):
train_transform = Compose([
RandomResizedCrop(width=512,
height=512,
scale=(0.08, 1.0),
ratio=(0.75, 1.3333333333333333)),
Transpose(p=0.5),
HorizontalFlip(p=0.5),
VerticalFlip(p=0.5),
ShiftScaleRotate(p=0.5),
HueSaturationValue(p=0.5),
RandomBrightnessContrast(brightness_limit=(-0.1, 0.1),
contrast_limit=(-0.1, 0.1),
p=0.5),
Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225],
max_pixel_value=255.0,
p=1.0),
CoarseDropout(p=0.5),
ToTensorV2(p=1.0)
],
p=1)
return train_transform
def __init__(self, target_size):
transformation = []
transformation += [Resize(target_size, target_size)]
transformation += [
RandomResizedCrop(target_size, target_size, scale=(0.8, 1.0)),
OneOf([
ShiftScaleRotate(),
GridDistortion(),
OpticalDistortion(),
ElasticTransform(approximate=True)
],
p=0.3),
OneOf([
IAAAdditiveGaussianNoise(),
GaussNoise(),
MedianBlur(blur_limit=3),
Blur(blur_limit=3)
],
p=0.3)
]
transformation += [Normalize(), ToTensor()]
print(transformation)
self.transform = Compose(transformation)
def generate_transforms_vit(img_size):
train_transform = Compose([
RandomResizedCrop(img_size, img_size),
Transpose(p=0.5),
HorizontalFlip(p=0.5),
VerticalFlip(p=0.5),
ShiftScaleRotate(p=0.5),
HueSaturationValue(hue_shift_limit=0.2,
sat_shift_limit=0.2,
val_shift_limit=0.2,
p=0.5),
RandomBrightnessContrast(
brightness_limit=(-0.1, 0.1), contrast_limit=(-0.1, 0.1), p=0.5),
Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225],
max_pixel_value=255.0,
p=1.0),
CoarseDropout(p=0.5),
Cutout(p=0.5),
ToTensorV2(p=1.0),
],
p=1.)
val_transform = Compose([
Resize(img_size, img_size),
HorizontalFlip(p=0.5),
VerticalFlip(p=0.5),
Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225],
max_pixel_value=255.0,
p=1.0),
ToTensorV2(p=1.0),
],
p=1.)
return {"train": train_transform, "val": val_transform}
def get_train_transforms(img_size):
return Compose([
RandomResizedCrop(img_size, img_size),
#RandomCrop(224, 224),
OneOf([
RandomGamma(gamma_limit=(60, 120), p=0.9),
RandomBrightnessContrast(brightness_limit=0.2,
contrast_limit=0.2,
p=0.9),
CLAHE(clip_limit=4.0, tile_grid_size=(4, 4), p=0.9),
]),
OneOf([
Blur(blur_limit=3, p=1),
MotionBlur(blur_limit=3, p=1),
MedianBlur(blur_limit=3, p=1)
],
p=0.5),
HorizontalFlip(p=0.5),
VerticalFlip(p=0.5),
HueSaturationValue(hue_shift_limit=0.2,
sat_shift_limit=0.2,
val_shift_limit=0.2,
p=0.5),
ShiftScaleRotate(shift_limit=0.2,
scale_limit=0.2,
rotate_limit=20,
interpolation=cv2.INTER_LINEAR,
border_mode=cv2.BORDER_CONSTANT,
p=1),
Normalize(mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225),
max_pixel_value=255.0,
p=1.0),
CoarseDropout(p=0.5),
ToTensorV2(p=1.0),
])
def main():
args = get_args()
os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda_ids
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
if not os.path.exists(args.save_path):
os.makedirs(args.save_path)
if not os.path.exists(args.tensorboard_dir):
os.makedirs(args.tensorboard_dir)
logging = get_logging()
seed_everything(args.seed)
transforms_train = Compose([
RandomResizedCrop(args.img_size, args.img_size),
Transpose(p=0.5),
HorizontalFlip(p=0.5),
VerticalFlip(p=0.5),
ShiftScaleRotate(p=0.5),
HueSaturationValue(hue_shift_limit=0.2, sat_shift_limit=0.2, val_shift_limit=0.2, p=0.5),
RandomBrightnessContrast(brightness_limit=(-0.1, 0.1), contrast_limit=(-0.1, 0.1), p=0.5),
Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], max_pixel_value=255.0, p=1.0),
CoarseDropout(p=0.5),
ToTensorV2(p=1.0),
], p=1.)
transforms_val = Compose([
CenterCrop(args.img_size, args.img_size, p=1.),
Resize(args.img_size, args.img_size),
Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], max_pixel_value=255.0, p=1.0),
ToTensorV2(p=1.0),
], p=1.)
train_loader, val_loader = get_dataset(args.data_path, args.batch_size, args.batch_size_val, transforms_train,
transforms_val)
net = get_model(True, device)
trainer = Trainer(net, train_loader, val_loader, args, device, logging)
trainer.train()
def get_transform(self, is_train=True):
if is_train:
return Compose([
RandomResizedCrop(self.config.image_size,
self.config.image_size),
Transpose(p=0.5),
HorizontalFlip(p=0.5),
VerticalFlip(p=0.5),
ShiftScaleRotate(p=0.5),
Normalize(
mean=self.settings.MEAN,
std=self.settings.STD,
),
ToTensorV2(),
])
else:
return Compose([
Resize(self.config.image_size, self.config.image_size),
Normalize(
mean=self.settings.MEAN,
std=self.settings.STD,
),
ToTensorV2(),
])
def __init__(self):
self.transformer = Compose([
HorizontalFlip(p=0.5),
CLAHE(clip_limit=(1, 8), tile_grid_size=(10, 10), p=0.3),
OneOf([
GridDistortion(num_steps=5, distort_limit=(-0.46, 0.40)),
ElasticTransform(alpha=1.68, sigma=48.32, alpha_affine=44.97),
],
p=0.3),
RandomResizedCrop(p=0.3,
height=512,
width=512,
scale=(0.08, 1.0),
ratio=(0.75, 1.33)),
ShiftScaleRotate(p=0.3,
shift_limit=(-0.06, 0.06),
scale_limit=(-0.10, 0.10),
rotate_limit=(-20, 20)),
Normalize(mean=[0.46009655, 0.43957878, 0.41827092],
std=[0.2108204, 0.20766491, 0.21656131],
max_pixel_value=255.0,
p=1.0),
ToTensorV2(),
])
def __init__(self, num_classes, input_size):
super(TuSimpleDataTransform, self).__init__()
random.seed(1000)
height, width = input_size
self._train_transform_list = self._train_transform_list + [
HorizontalFlip(p=0.5),
GaussNoise(p=0.5),
RandomBrightnessContrast(p=0.5),
RandomShadow(p=0.5),
RandomRain(rain_type="drizzle", p=0.5),
ShiftScaleRotate(rotate_limit=10, p=0.5),
RandomResizedCrop(
height=height, width=width, scale=(0.8, 1), p=0.5),
]
self._train_transform_list.append(Resize(height, width))
self._val_transform_list.append(Resize(height, width))
self._train_transform_list.append(ToTensor(num_classes=num_classes))
self._val_transform_list.append(ToTensor(num_classes=num_classes))
self._initialize_transform_dict()
def get_inference_transforms(input_shape, way="pad", crop_rate=1.0):
if way == "pad":
return Compose(
[
PadIfNeeded(input_shape[0], input_shape[1]),
Resize(input_shape[0], input_shape[1]),
HorizontalFlip(p=0.5),
ToGray(p=0.5),
VerticalFlip(p=0.5),
ShiftScaleRotate(scale_limit=0.0, p=0.5),
HueSaturationValue(hue_shift_limit=0.2,
sat_shift_limit=0.2,
val_shift_limit=0.2,
p=0.5),
RandomBrightnessContrast(brightness_limit=(-0.1, 0.1),
contrast_limit=(-0.1, 0.1),
p=0.5),
Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225],
max_pixel_value=255.0,
p=1.0),
#CoarseDropout(p=0.5),
Cutout(p=0.5),
ToTensorV2(p=1.0),
],
p=1.)
elif way == "resize":
return Compose(
[
RandomResizedCrop(input_shape[0], input_shape[1]),
HorizontalFlip(p=0.5),
ToGray(p=0.5),
VerticalFlip(p=0.5),
ShiftScaleRotate(scale_limit=0.0, p=0.5),
HueSaturationValue(hue_shift_limit=0.2,
sat_shift_limit=0.2,
val_shift_limit=0.2,
p=0.5),
RandomBrightnessContrast(brightness_limit=(-0.1, 0.1),
contrast_limit=(-0.1, 0.1),
p=0.5),
Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225],
max_pixel_value=255.0,
p=1.0),
#CoarseDropout(p=0.5),
Cutout(p=0.5),
ToTensorV2(p=1.0),
],
p=1.)
elif way == "center":
return Compose(
[
Resize(input_shape[0], input_shape[1]),
HorizontalFlip(p=0.5),
ToGray(p=0.5),
VerticalFlip(p=0.5),
ShiftScaleRotate(scale_limit=0.0, p=0.5),
HueSaturationValue(hue_shift_limit=0.2,
sat_shift_limit=0.2,
val_shift_limit=0.2,
p=0.5),
RandomBrightnessContrast(brightness_limit=(-0.1, 0.1),
contrast_limit=(-0.1, 0.1),
p=0.5),
Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225],
max_pixel_value=255.0,
p=1.0),
#CoarseDropout(p=0.5),
Cutout(p=0.5),
ToTensorV2(p=1.0),
],
p=1.)
elif way == "crop":
return Compose(
[
Resize(input_shape[0], input_shape[1]),
CenterCrop(int(input_shape[0] * crop_rate),
int(input_shape[1] * crop_rate)),
HorizontalFlip(p=0.5),
ToGray(p=0.5),
VerticalFlip(p=0.5),
ShiftScaleRotate(scale_limit=0.0, p=0.5),
HueSaturationValue(hue_shift_limit=0.2,
sat_shift_limit=0.2,
val_shift_limit=0.2,
p=0.5),
RandomBrightnessContrast(brightness_limit=(-0.1, 0.1),
contrast_limit=(-0.1, 0.1),
p=0.5),
Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225],
max_pixel_value=255.0,
p=1.0),
#CoarseDropout(p=0.5),
Cutout(p=0.5),
ToTensorV2(p=1.0),
],
p=1.)
hvd.init()
print("\n********using distributed training:%d***********" % hvd.rank())
torch.cuda.set_device(hvd.local_rank())
# Horovod: limit # of CPU threads to be used per worker.
torch.set_num_threads(4)
if (not args.distribute) or (args.distribute and hvd.rank() == 0):
print(args)
output_dir = args.output
os.makedirs(output_dir + '/checkpoint', exist_ok=True)
mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
# print(model)
train_augment = Compose([
RandomResizedCrop(224, 224),
HorizontalFlip(p=0.5),
Normalize(mean, std),
ToTensorV2(),
])
val_augment = Compose([
Resize(256, 256),
CenterCrop(224, 224),
Normalize(mean, std),
ToTensorV2(),
])
best_val_score = [0, 0, 0, 0, 0]
for fold in [0]: # [0, 1, 2, 3, 4]
if (not args.distribute) or (args.distribute and hvd.rank() == 0):
print(
from torch.utils.data import Dataset, Sampler, DataLoader
from tqdm import tqdm
ORIG_SHAPE = (1400, 2100)
TRAIN_SHAPE = (352, 512)
NUM_CLASSES = 4
bad_images = [
'046586a', '1588d4c', '1e40a05', '41f92e5', '449b792', '563fc48',
'8bd81ce', 'b092cc1', 'c0306e5', 'c26c635', 'e04fea3', 'e5f2f24',
'eda52f2', 'fa645da'
]
AUGMENTATIONS_TRAIN = Compose([
RandomRotate90(p=0.25),
RandomResizedCrop(
height=TRAIN_SHAPE[0], width=TRAIN_SHAPE[1], scale=(0.8, 1.0)),
HorizontalFlip(p=0.5),
VerticalFlip(p=0.5),
ShiftScaleRotate(shift_limit=(-0.2, 0.2),
scale_limit=(-0.2, 0.2),
rotate_limit=(-20, 20),
border_mode=0,
interpolation=1,
p=0.25),
OneOf([
RandomBrightnessContrast(brightness_limit=(-0.2, 0.2),
contrast_limit=(-0.2, 0.2)),
RandomGamma(),
CLAHE()
],
p=0.4),
def train(model, cfg, model_cfg, start_epoch=0):
cfg.batch_size = 16 if cfg.batch_size < 1 else cfg.batch_size
cfg.val_batch_size = cfg.batch_size
cfg.input_normalization = model_cfg.input_normalization
crop_size = model_cfg.crop_size
loss_cfg = edict()
loss_cfg.pixel_loss = MaskWeightedMSE(min_area=100)
loss_cfg.pixel_loss_weight = 1.0
num_epochs = 180
train_augmentator = HCompose(
[RandomResizedCrop(*crop_size, scale=(0.5, 1.0)),
HorizontalFlip()])
val_augmentator = HCompose([Resize(*crop_size)])
trainset = ComposeDataset(
[
HDataset(cfg.HFLICKR_PATH, split='train'),
HDataset(cfg.HDAY2NIGHT_PATH, split='train'),
HDataset(cfg.HCOCO_PATH, split='train'),
HDataset(cfg.HADOBE5K_PATH, split='train'),
],
augmentator=train_augmentator,
input_transform=model_cfg.input_transform,
keep_background_prob=0.05,
)
valset = ComposeDataset(
[
HDataset(cfg.HFLICKR_PATH, split='test'),
HDataset(cfg.HDAY2NIGHT_PATH, split='test'),
HDataset(cfg.HCOCO_PATH, split='test'),
],
augmentator=val_augmentator,
input_transform=model_cfg.input_transform,
keep_background_prob=-1,
)
optimizer_params = {'lr': 1e-3, 'betas': (0.9, 0.999), 'eps': 1e-8}
lr_scheduler = partial(torch.optim.lr_scheduler.MultiStepLR,
milestones=[160, 175],
gamma=0.1)
trainer = SimpleHTrainer(
model,
cfg,
model_cfg,
loss_cfg,
trainset,
valset,
optimizer='adam',
optimizer_params=optimizer_params,
lr_scheduler=lr_scheduler,
metrics=[
DenormalizedPSNRMetric(
'images',
'target_images',
mean=torch.tensor(cfg.input_normalization['mean'],
dtype=torch.float32).view(1, 3, 1, 1),
std=torch.tensor(cfg.input_normalization['std'],
dtype=torch.float32).view(1, 3, 1, 1),
),
DenormalizedMSEMetric(
'images',
'target_images',
mean=torch.tensor(cfg.input_normalization['mean'],
dtype=torch.float32).view(1, 3, 1, 1),
std=torch.tensor(cfg.input_normalization['std'],
dtype=torch.float32).view(1, 3, 1, 1),
)
],
checkpoint_interval=10,
image_dump_interval=1000)
logger.info(f'Starting Epoch: {start_epoch}')
logger.info(f'Total Epochs: {num_epochs}')
for epoch in range(start_epoch, num_epochs):
trainer.training(epoch)
trainer.validation(epoch)
def generate_transforms(image_size):
# train_transform = Compose(
# [
# Resize(height=int(image_size[0]), width=int(image_size[1])),
# OneOf([RandomBrightness(limit=0.1, p=1), RandomContrast(limit=0.1, p=1)]),
# OneOf([MotionBlur(blur_limit=3), MedianBlur(blur_limit=3), GaussianBlur(blur_limit=3)], p=0.5),
# VerticalFlip(p=0.5),
# HorizontalFlip(p=0.5),
# # Transpose(p=0.5),
# HueSaturationValue(hue_shift_limit=0.2, sat_shift_limit=0.2, val_shift_limit=0.2, p=0.5),
# CoarseDropout(p=0.5),
# Cutout(p=0.5),
# ShiftScaleRotate(
# shift_limit=0.2,
# scale_limit=0.2,
# rotate_limit=20,
# interpolation=cv2.INTER_LINEAR,
# border_mode=cv2.BORDER_REFLECT_101,
# p=1,
# ),
# Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0, p=1.0),
# ]
# )
train_transform = Compose(
[
RandomResizedCrop(int(image_size[0]),
int(image_size[1]),
scale=(0.08, 1.0),
ratio=(0.75, 1.3333333333333333)),
Transpose(p=0.5),
HorizontalFlip(p=0.5),
VerticalFlip(p=0.5),
ShiftScaleRotate(p=0.5),
HueSaturationValue(hue_shift_limit=0.2,
sat_shift_limit=0.2,
val_shift_limit=0.2,
p=0.5),
RandomBrightnessContrast(brightness_limit=(-0.1, 0.1),
contrast_limit=(-0.1, 0.1),
p=0.5),
Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225],
max_pixel_value=255.0,
p=1.0),
# CoarseDropout(p=0.5),
# Cutout(p=0.5),
# ToTensorV2(p=1.0),
],
p=1.)
val_transform = Compose([
Resize(height=int(image_size[0]), width=int(image_size[1])),
Normalize(mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225),
max_pixel_value=255.0,
p=1.0),
])
return {
"train_transforms": train_transform,
"val_transforms": val_transform
}
transformed_image = transform(image=image)
elif augmentation == 'center_crop_to_square':
transform = iaa.CenterCropToSquare()
transformed_image = transform(image=image)
elif augmentation == 'random_crop':
transform = RandomCrop(always_apply=True, width=200, height=200)
transformed_image = transform(image=image)['image']
elif augmentation == 'crop_and_pad':
transform = iaa.CropAndPad(percent=(-0.25, 0.25))
transformed_image = transform(image=image)
elif augmentation == 'random_resized_crop':
transform = RandomResizedCrop(always_apply=True, width=100, height=100)
transformed_image = transform(image=image)['image']
elif augmentation == 'random_sized_crop':
transform = RandomSizedCrop(always_apply=True, height=500, width=500,
min_max_height=[200, 200])
transformed_image = transform(image=image)['image']
## Distortion
elif augmentation == 'grid_distortion':
transform = GridDistortion(always_apply=True, distort_limit=0.5)
transformed_image = transform(image=image)['image']
elif augmentation == 'optical_distortion':
transform = OpticalDistortion(always_apply=True, distort_limit=0.5)
def __init__(self,
split='train',
train_ratio=0.95,
useAug=True,
inputtrans="",
gttrans="",
fast_load=True):
"""
train_ratio: ratio of train
1e4 pairs in all
"""
assert inputtrans
assert gttrans
# self.input_dir = "/data/ntire/rawdata/train/input/"
self.input_dir = TR_INPUT
# ex: 000006_3.png
# self.gt_dir = "/data/ntire/rawdata/train/gt/"
self.gt_dir = TR_GT
# ex: 000003_gt.png
self.input_list = glob.glob(self.input_dir + '/*.png')
self.gt_list = glob.glob(self.gt_dir + '/*.png')
self.input_transform = get_trans(inputtrans)
self.gt_transform = get_trans(gttrans)
self.input_list = np.array(sorted(self.input_list))
self.gt_list = np.array(sorted(self.gt_list))
self.split = split
self.len = len(self.input_list)
self.useAug = useAug
if (self.split == 'train'):
with open('train.pkl', 'rb') as ff:
train_ids = pkl.load(ff)
self.input_list = self.input_list[train_ids]
self.gt_list = self.gt_list[train_ids]
# aug
if (useAug):
self.aug = A.Compose([
RandomRotate90(p=0.1),
Flip(p=0.1),
Transpose(p=0.1),
RandomResizedCrop(height=128, width=128, p=0.1)
],
additional_targets={
'image': 'image',
'gt': 'image'
})
elif (self.split == 'val'):
with open('val.pkl', 'rb') as ff:
val_ids = pkl.load(ff)
self.input_list = self.input_list[val_ids]
self.gt_list = self.gt_list[val_ids]
self.fast_load = fast_load
if self.fast_load:
self.input_images = {}
idx = [i for i in range(len(self.input_list))]
target = ['self.input_images'] * len(self.input_list)
inp = zip(idx, self.input_list, target)
from tqdm import tqdm
for i in tqdm(list(inp), desc='loading input images'):
self.save_in_memory(i)
# Pool didn't work ????
# with Pool(5) as p:
# p.map(self.save_in_memory, list(inp))
self.gt_images = {}
idx = [i for i in range(len(self.gt_list))]
target = ['self.gt_images'] * len(self.gt_list)
inp = zip(idx, self.gt_list, target)
for i in tqdm(list(inp), desc='loading gt images'):
self.save_in_memory(i)
# with Pool(5) as p:
# p.map(self.save_in_memory, list(inp))
self.len = len(self.input_list)
self.totensor = ToTensor()
logger.info(f'Loaded Dataset split:{split}, len:{self.len}')
},
'train': {
'mean': [0.485, 0.456, 0.406],
'std': [0.229, 0.224, 0.225],
'loss': EXPERIMENT_TYPE_LOSS_SIMPLE,
'data': EXPERIMENT_TYPE_USE_STEROID,
'init_lr': 0.0005,
'lee': 0,
'scaling': True,
'contact': 13,
'batch_size': 12,
'epochs': 100,
'train_transform': Compose(
[
OneOf([RandomCrop(512, 512, p=0.7),
RandomResizedCrop(512, 512, scale=(0.6, 1.0), p=0.3)],
p=1.0),
RandomBrightnessContrast(p=0.3),
Blur(p=0.3),
Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
ToTensor(),
]),
'valid_transform': Compose(
[CenterCrop(512, 512, p=1.0),
Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
ToTensor()
]),
'infer_transform': Compose(
from albumentations import RandomResizedCrop, Compose, Flip, Resize, Normalize
from albumentations.pytorch import ToTensorV2, ToTensor
SIZE = 224
train_aug = Compose([
RandomResizedCrop(height=SIZE, width=SIZE, ),
Flip(p=0.3),
Normalize(),
ToTensorV2(),
])
valid_aug = Compose([
Resize(width=SIZE, height=SIZE),
Normalize(),
ToTensorV2(),
])
def train(**kwargs):
parser = argparse.ArgumentParser()
parser.add_argument("--config", default="./vanilla.json", type=str,
help="path to the config file")
args = parser.parse_args()
args = Config(args.config)
get_dataset_with_arg = {"tiny_imagenet": food.datasets.TinyImagenet,
"cifar_100": food.datasets.CIFAR_100,
"cifar_10": food.datasets.CIFAR_10}
if not args.keep_logs:
try:
shutil.rmtree(args.logdir)
except FileNotFoundError:
pass
batch_size = kwargs.get('batch_size', args.batch_size)
model = kwargs.get('model', args.model).lower()
dataset = kwargs.get('dataset', args.dataset).lower()
epochs = kwargs.get('epochs', args.epochs)
alpha = kwargs.get('alpha', args.alpha)
loss_type = kwargs.get('loss_type', args.loss_type).lower()
test_b = kwargs.get('test', False)
ds_class = get_dataset_with_arg[dataset.lower()]
train_transforms = Compose([
HorizontalFlip(p=0.5),
RandomResizedCrop(32, 32, p=0.5),
Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
ToTensor(),
], p=1)
val_transforms = Compose([
Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
ToTensor(),
], p=1)
train_dataset = ds_class(args.data_path, mode="train", task=args.task.lower(),
transform=train_transforms)
val_dataset = ds_class(args.data_path, mode="val", task=args.task.lower(),
transform=val_transforms)
model = getattr(torchvision.models, args.model)(
num_classes=train_dataset.n_classes)
model = Net2LastLayers(model)
print("Total number of model's parameters: ",
np.sum([p.numel() for p in model.parameters() if p.requires_grad]))
train_dataloader = DataLoader(
train_dataset, batch_size=batch_size, shuffle=True, drop_last=True)
val_dataloader = DataLoader(
val_dataset, batch_size=batch_size, shuffle=False, drop_last=False)
# if torch.cuda.is_available(): TODO
# train_dataloader = DataPrefetcher(train_dataloader)
# val_dataloader = DataPrefetcher(val_dataloader)
# pass
device = "cuda:0" if torch.cuda.is_available() else "cpu"
model.to(device)
criterion = GlassLoss(alpha=alpha, loss_type=loss_type)
optimizer = torch.optim.Adam(lr=args.lr, params=model.parameters())
scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer, [50, 100, 200], gamma=0.1) # ExponentialLR(optimizer, 0.9)
if args.resume is not None:
state_dict = torch.load(args.resume, map_location=device)
model.load_state_dict(state_dict["model_state_dict"])
optimizer.load_state_dict(state_dict["optimizer_state_dict"])
train_writer = SummaryWriter(os.path.join(args.logdir, "train_logs"))
val_writer = SummaryWriter(os.path.join(args.logdir, "val_logs"))
global_step = 0
for epoch in range(epochs):
print(f"Training, epoch {epoch + 1}")
model.train()
for images, labels in train_dataloader:
images, labels = images.to(device), labels.to(device)
optimizer.zero_grad()
prev_logits, logits, V = model.forward(images)
prev_logits, logits, V = prev_logits.to(
device), logits.to(device), V.to(device)
# print(V)
# print(logits.shape)
# print(prev_logits.shape)
loss = criterion(model, V.to(device), labels, prev_logits, logits)
predictions = logits.argmax(dim=1)
accuracy_t = torch.mean((predictions == labels).float()).item()
if global_step % args.log_each == 0:
# TODO
# train_writer.add_scalar("Loss_BCE", loss, global_step=global_step)
# train_writer.add_scalar("Accuracy", accuracy_t, global_step=global_step)
# train_writer.add_scalar("Learning_rate", scheduler.get_lr()[-1], global_step=global_step)
#log_dict_with_writer(labels, logits, train_writer, global_step=global_step)
pass
loss.backward()
optimizer.step()
global_step += 1
print("Validating...")
val_loss, val_acc = evaluate(
model, val_dataloader, criterion, device, val_writer)
val_writer.add_scalar("Loss_BCE", val_loss, global_step=global_step)
val_writer.add_scalar("Accuracy", val_acc, global_step=global_step)
if epoch % args.checkpoint_each == 0:
print("Saving checkpoint...")
with open(os.path.join(args.checkpoints_dir, f"epoch_{epoch}_{global_step}.pt"), "wb") as f:
torch.save({"model_state_dict": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict()}, f)
scheduler.step()
if test_b:
return logits, loss, predictions, val_loss, val_acc
# RandomSizedCrop
aug = RandomSizedCrop(p=1, min_max_height=(100, 200), height=128, width=128)
augmented = aug(image=image, mask=mask)
image_scaled = augmented["image"]
mask_scaled = augmented["mask"]
visualize(image_scaled, mask_scaled, original_image=image, original_mask=mask)
# %%
# RandomResizedCrop
aug = RandomResizedCrop(p=1, height=72, width=72, scale=(0.25, 1.0))
augmented = aug(image=image, mask=mask)
image_scaled = augmented["image"]
mask_scaled = augmented["mask"]
visualize(image_scaled, mask_scaled, original_image=image, original_mask=mask)
# %%
# CLAHE
aug = CLAHE()
image8 = (image * 256).astype("uint8")
augmented = aug(image=image8, mask=mask.astype("uint8"))
