def create_transform(self, args, is_train):
"""
Convert numpy array into Tensor if dataset is for validation.
Apply data augmentation method to train dataset while cv or test if args.use_aug is 1.
is_train: boolean
flg that dataset is for validation in cv or test
return: Compose of albumentations
"""
if is_train and args.use_aug == 1:
transform = A.Compose([
trans.Resize(299, 299),
trans.Normalize(mean=self.img_mean,
std=self.img_std,
max_pixel_value=1.0),
ToTensorV2()
])
else:
transform = A.Compose([
trans.Resize(299, 299),
trans.Normalize(mean=self.img_mean,
std=self.img_std,
max_pixel_value=1.0),
ToTensorV2()
])
return transform
def augmentation(mode, target_size, prob=0.5, aug_m=2):
'''
description: augmentation
mode: 'train' 'test'
target_size: int or list, the shape of image ,
aug_m: Strength of transform
'''
high_p = prob
low_p = high_p / 2.0
M = aug_m
first_size = [int(x / 0.7) for x in target_size]
if mode == 'train':
return composition.Compose([
transforms.Resize(first_size[0], first_size[1], interpolation=3),
transforms.Flip(p=0.5),
composition.OneOf([
RandomCenterCut(scale=0.1 * M),
transforms.ShiftScaleRotate(shift_limit=0.05 * M,
scale_limit=0.1 * M,
rotate_limit=180,
border_mode=cv2.BORDER_CONSTANT,
value=0),
albumentations.imgaug.transforms.IAAAffine(
shear=(-10 * M, 10 * M), mode='constant')
],
p=high_p),
transforms.RandomBrightnessContrast(
brightness_limit=0.1 * M, contrast_limit=0.03 * M, p=high_p),
transforms.HueSaturationValue(hue_shift_limit=5 * M,
sat_shift_limit=15 * M,
val_shift_limit=10 * M,
p=high_p),
transforms.OpticalDistortion(distort_limit=0.03 * M,
shift_limit=0,
border_mode=cv2.BORDER_CONSTANT,
value=0,
p=low_p),
composition.OneOf([
transforms.Blur(blur_limit=7),
albumentations.imgaug.transforms.IAASharpen(),
transforms.GaussNoise(var_limit=(2.0, 10.0), mean=0),
transforms.ISONoise()
],
p=low_p),
transforms.Resize(target_size[0], target_size[1], interpolation=3),
transforms.Normalize(mean=(0.5, 0.5, 0.5),
std=(0.5, 0.5, 0.5),
max_pixel_value=255.0)
],
p=1)
else:
return composition.Compose([
transforms.Resize(target_size[0], target_size[1], interpolation=3),
transforms.Normalize(mean=(0.5, 0.5, 0.5),
std=(0.5, 0.5, 0.5),
max_pixel_value=255.0)
],
p=1)
def train_transform(resize, normalize=None):
if normalize == 'imagenet':
trans_fucn = [
albu.VerticalFlip(p=0.5),
albu.HorizontalFlip(p=0.5),
# albu.ToFloat(max_value=255, p=1.0),
albu.Normalize(p=1.0),
ToTensorV2(p=1.0)
]
elif normalize == 'global_norm':
trans_fucn = [
albu.VerticalFlip(p=0.5),
albu.HorizontalFlip(p=0.5),
GlobalNormalize(p=1.0),
# albu.ToFloat(max_value=255, p=1.0),
ToTensorV2(p=1.0)
]
else:
trans_fucn = [
albu.VerticalFlip(p=0.5),
albu.HorizontalFlip(p=0.5),
albu.Normalize(mean=(0, 0, 0), std=(1, 1, 1)),
# albu.ToFloat(max_value=255, p=1.0),
ToTensorV2(p=1.0)
]
return Compose(trans_fucn, p=1.0)
def get_transformv2(opt):
transform_list = []
# Transforms in opt.preprocess
if 'fixsize' in opt.preprocess:
transform_list.append(tr.Resize(286, 286, interpolation=2, p=1))
if 'resize' in opt.preprocess:
transform_list.append(
tr.Resize(opt.load_size, opt.load_size, interpolation=2, p=1))
if 'crop' in opt.preprocess:
transform_list.append(tr.RandomCrop(opt.crop_size, opt.crop_size, p=1))
# Transforms in colorspace
if 'color' in opt.preprocess:
transform_list.extend([
tr.RandomContrast(limit=0.2, p=0.5),
tr.RandomBrightness(limit=0.2, p=0.5),
tr.HueSaturationValue(hue_shift_limit=20,
sat_shift_limit=30,
val_shift_limit=20,
p=0.5),
# tr.ISONoise()
])
# Necessary transforms
transform_list.extend([
tr.HorizontalFlip(p=0.5),
tr.VerticalFlip(p=0.5),
tr.Normalize(p=1.0),
ToTensorV2(p=1)
])
return Compose(transform_list, additional_targets={'imageB': 'image'})
def ben_valid_augmentation():
return ACompose([
atransforms.Resize(128, 128, interpolation=3),
atransforms.Normalize(mean=BEN_BAND_STATS['mean'],
std=BEN_BAND_STATS['std']),
AToTensor(),
])
def eval_transform(resize, normalize=None):
if normalize == 'imagenet':
trans_func = [
albu.Normalize(p=1.0),
ToTensorV2(p=1.0)
]
elif normalize == 'global_norm':
trans_func = [
GlobalNormalize(p=1.0),
ToTensorV2(p=1.0)
]
else:
trans_func = [
albu.Normalize(mean=(0, 0, 0), std=(1, 1, 1)),
ToTensorV2(p=1.0)
]
return Compose(trans_func, p=1.0)
def get_presize_combine_transforms_V4():
transforms_presize = A.Compose([
transforms.PadIfNeeded(600, 800),
geometric.Perspective(
scale=[0, .1],
pad_mode=cv2.BORDER_REFLECT,
interpolation=cv2.INTER_AREA, p = .3),
transforms.Flip(),
geometric.ShiftScaleRotate(interpolation=cv2.INTER_LANCZOS4, p = 0.95, scale_limit=0.0),
crops.RandomResizedCrop(
TARGET_SIZE, TARGET_SIZE,
scale=(config['rrc_scale_min'], config['rrc_scale_max']),
ratio=(.70, 1.4),
interpolation=cv2.INTER_CUBIC,
p=1.0),
transforms.Transpose()
#rotate.Rotate(interpolation=cv2.INTER_LANCZOS4, p = 0.99),
])
transforms_postsize = A.Compose([
#imgaug.IAAPiecewiseAffine(),
transforms.CoarseDropout(),
transforms.CLAHE(p=.1),
transforms.RandomToneCurve(scale=.1, p=0.2),
transforms.RandomBrightnessContrast(
brightness_limit=.1,
contrast_limit=0.4,
p=.8),
transforms.HueSaturationValue(
hue_shift_limit=20,
sat_shift_limit=50,
val_shift_limit=0,
p=0.5),
transforms.Equalize(p=0.05),
transforms.FancyPCA(p=0.05),
transforms.RandomGridShuffle(p=0.1),
A.OneOf([
transforms.MotionBlur(blur_limit=(3, 9)),
transforms.GaussianBlur(),
transforms.MedianBlur()
], p=0.1),
transforms.ISONoise(p=.2),
transforms.GaussNoise(var_limit=127., p=.3),
A.OneOf([
transforms.GridDistortion(interpolation=cv2.INTER_AREA, distort_limit=[0.7, 0.7], p=0.5),
transforms.OpticalDistortion(interpolation=cv2.INTER_AREA, p=.3),
], p=.3),
geometric.ElasticTransform(alpha=4, sigma=4, alpha_affine=4, interpolation=cv2.INTER_AREA, p=0.3),
transforms.CoarseDropout(),
transforms.Normalize(),
ToTensorV2()
])
return transforms_presize, transforms_postsize
def AlbumentationTrainTransform(self):
tf = tc.Compose([ta.HorizontalFlip(p=0.5),
ta.Rotate(limit=(-20, 20)),
# ta.VerticalFlip(p=0.5),
# ta.Cutout(num_holes=3, max_h_size=8, max_w_size=8, p=0.5),
# ta.Blur(),
# ta.ChannelShuffle(),
# ta.InvertImg(),
ta.RandomCrop(height=30, width=30, p=5.0),
ta.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
tp.ToTensor()
])
return lambda img: tf(image = np.array(img))["image"]
def __init__(self, im_paths=None, labels=None, phase=None, resize=False):
"""
Args:
im_paths (numpy): image_data
y (numpy): label data
transform: pytorch transforms for transforms and tensor conversion
"""
self.im_paths = im_paths
self.labels = labels
self.resize = resize
self.albumentations_transform = {
'train':
Compose([
ab_transforms.HorizontalFlip(p=0.2),
ab_transforms.VerticalFlip(p=0.2),
ab_transforms.Rotate(limit=180, p=0.2),
ab_transforms.HueSaturationValue(p=0.1),
ab_transforms.RandomContrast(p=0.1),
ab_transforms.GaussianBlur(blur_limit=3, p=0.2),
ab_transforms.GaussNoise(p=0.05),
ab_transforms.CLAHE(p=0.2),
ab_transforms.Normalize(mean=[0.5944, 0.4343, 0.5652],
std=[0.2593, 0.2293, 0.2377]),
ToTensor()
]),
'val':
Compose([
ab_transforms.Normalize(mean=[0.5944, 0.4343, 0.5652],
std=[0.2593, 0.2293, 0.2377]),
ToTensor()
]),
}
if phase == 'train':
self.transform = self.albumentations_transform['train']
else:
self.transform = self.albumentations_transform['val']
def create_transform(self, args, is_train):
"""
Convert numpy array into Tensor if dataset is for validation.
Apply data augmentation method to train dataset while cv or test if args.use_aug is 1.
is_train: boolean
flg that dataset is for validation in cv or test
return: Compose of albumentations
"""
if is_train and args.use_aug == 1:
transform = A.Compose([
trans.Normalize(mean=self.cifar_10_mean,
std=self.cifar_10_std,
max_pixel_value=1.0),
trans.HorizontalFlip(p=0.5),
trans.ShiftScaleRotate(shift_limit=0,
scale_limit=0.25,
rotate_limit=30,
p=1),
trans.CoarseDropout(max_holes=1,
min_holes=1,
min_width=12,
min_height=12,
max_height=12,
max_width=12,
p=0.5),
ToTensorV2()
])
else:
transform = A.Compose([
trans.Normalize(mean=self.cifar_10_mean,
std=self.cifar_10_std,
max_pixel_value=1.0),
ToTensorV2()
])
return transform
def __init__(self, data_folder, cli_args):
self.cli_args = cli_args
self.root: str = data_folder
self.image_names: List[str] = sorted(os.listdir(os.path.join(self.root,
"test",
"imgs")))
self.transform = Compose(
[
# Normalize images to [0..1]
tf.Normalize(mean=(0.0, 0.0, 0.0), std=(1.0, 1.0, 1.0), p=1),
# Resize images to (image_size, image_size)
tf.Resize(cli_args.image_size, cli_args.image_size),
# Convert PIL images to torch.Tensor
ToTensorV2(),
]
)
def get_presize_combine_tune_transforms():
transforms_presize = A.Compose([
transforms.Transpose(),
transforms.Flip(),
#transforms.PadIfNeeded(600, 800),
crops.RandomResizedCrop(
TARGET_SIZE, TARGET_SIZE,
scale=(.75, 1),
interpolation=cv2.INTER_CUBIC,
p=1.0),
rotate.Rotate(interpolation=cv2.INTER_LANCZOS4, p = 0.99),
])
transforms_postsize = A.Compose([
transforms.CoarseDropout(),
# transforms.CLAHE(p=.1),
transforms.RandomToneCurve(scale=.1),
transforms.RandomBrightnessContrast(
brightness_limit=.1,
contrast_limit=0.2,
p=.7),
transforms.HueSaturationValue(
hue_shift_limit=20,
sat_shift_limit=60,
val_shift_limit=0,
p=0.6),
#transforms.Equalize(p=0.1),
#transforms.FancyPCA(p=0.05),
#transforms.RandomGridShuffle(p=0.1),
#A.OneOf([
# transforms.MotionBlur(blur_limit=(3, 9)),
# transforms.GaussianBlur(),
# transforms.MedianBlur()
# ], p=0.2),
transforms.ISONoise(p=.3),
transforms.GaussNoise(var_limit=255., p=.3),
#A.OneOf([
# transforms.GridDistortion(interpolation=cv2.INTER_AREA, distort_limit=[0.7, 0.7], p=0.5),
# transforms.OpticalDistortion(interpolation=cv2.INTER_AREA, p=.3),
# ], p=.3),
geometric.ElasticTransform(alpha=4, sigma=100, alpha_affine=100, interpolation=cv2.INTER_AREA, p=0.3),
transforms.CoarseDropout(),
transforms.Normalize(),
ToTensorV2()
])
return transforms_presize, transforms_postsize
def ben_augmentation():
# !! Need to do something to change color _probably_
return ACompose([
atransforms.HorizontalFlip(p=0.5),
atransforms.RandomRotate90(p=1.0),
atransforms.ShiftScaleRotate(shift_limit=0, scale_limit=0, p=1.0),
atransforms.RandomSizedCrop((60, 120),
height=128,
width=128,
interpolation=3),
# atransforms.GridDistortion(num_steps=5, p=0.5), # !! Maybe too much noise?
atransforms.Normalize(mean=BEN_BAND_STATS['mean'],
std=BEN_BAND_STATS['std']),
# atransforms.ChannelDropout(channel_drop_range=(1, 2), p=0.5),
AToTensor(),
])
def test_dataloader(self):
augmentations = Compose([
A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
ToTensorV2(),
])
test_ds = MelanomaDataset(
df=self.test_df,
images_path=self.test_images_path,
augmentations=augmentations, # TODO: add TTA
train_or_valid=False,
)
return DataLoader(
test_ds,
batch_size=self.hparams.bs,
shuffle=False,
num_workers=os.cpu_count(),
pin_memory=True,
)
def get_valid_transforms():
return A.Compose([
transforms.Transpose(),
transforms.PadIfNeeded(600, 800),
rotate.Rotate(interpolation=cv2.INTER_LANCZOS4, p = 0.90),
crops.RandomResizedCrop(
TARGET_SIZE_VALID, TARGET_SIZE_VALID,
scale=(.75, 1),
interpolation=cv2.INTER_CUBIC,
p=1.0),
transforms.Flip(),
#transforms.RandomToneCurve(scale=.1),
#transforms.RandomBrightnessContrast(brightness_limit=0.0, contrast_limit=0.3, p=.7),
#transforms.HueSaturationValue(hue_shift_limit=10,
# sat_shift_limit=10,
# val_shift_limit=5, p=0.6),
transforms.Normalize(),
ToTensorV2()
])
def get_train_transforms():
return A.Compose([
transforms.PadIfNeeded(600, 800),
geometric.ShiftScaleRotate(interpolation=cv2.INTER_LANCZOS4, p = 0.99, scale_limit=0.8),
geometric.Perspective(pad_mode=cv2.BORDER_REFLECT,interpolation=cv2.INTER_AREA),
crops.RandomResizedCrop(
TARGET_SIZE, TARGET_SIZE,
scale=(config['rrc_scale_min'], config['rrc_scale_max']),
interpolation=cv2.INTER_CUBIC,
p=1.0),
transforms.Transpose(),
transforms.Flip(),
transforms.CoarseDropout(),
transforms.CLAHE(p=.1),
transforms.RandomToneCurve(scale=.1),
transforms.RandomBrightnessContrast(
brightness_limit=.1,
contrast_limit=0.3,
p=.7),
transforms.HueSaturationValue(
hue_shift_limit=20,
sat_shift_limit=60,
val_shift_limit=0,
p=0.6),
transforms.RandomGridShuffle(p=0.1),
A.OneOf([
transforms.MotionBlur(blur_limit=(3, 9)),
transforms.GaussianBlur(),
transforms.MedianBlur()
], p=0.2),
transforms.ISONoise(p=.3),
transforms.GaussNoise(var_limit=255., p=.3),
A.OneOf([
transforms.GridDistortion(interpolation=cv2.INTER_AREA, distort_limit=[0.7, 0.7], p=0.5),
transforms.OpticalDistortion(interpolation=cv2.INTER_AREA, p=.3),
], p=.3),
geometric.ElasticTransform(alpha=4, sigma=100, alpha_affine=100, interpolation=cv2.INTER_AREA, p=0.3),
transforms.CoarseDropout(),
transforms.Normalize(),
ToTensorV2()
])
def train_dataloader(self):
augmentations = Compose(
[
A.RandomResizedCrop(
height=self.hparams.sz,
width=self.hparams.sz,
scale=(0.7, 1.0),
),
# AdvancedHairAugmentation(),
A.GridDistortion(),
A.RandomBrightnessContrast(),
A.ShiftScaleRotate(),
A.Flip(p=0.5),
A.CoarseDropout(
max_height=int(self.hparams.sz / 10),
max_width=int(self.hparams.sz / 10),
),
# A.HueSaturationValue(),
A.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225],
max_pixel_value=255,
),
ToTensorV2(),
]
)
train_ds = MelanomaDataset(
df=self.train_df,
images_path=self.train_images_path,
augmentations=augmentations,
train_or_valid=True,
)
return DataLoader(
train_ds,
# sampler=sampler,
batch_size=self.hparams.bs,
shuffle=True,
num_workers=os.cpu_count(),
pin_memory=True,
)
def get_patched_input(img_path, config, gt_mask_flag):
img_patch_set = []
p_size = config['patch_size']
img_size = config['input_w']
patch_overlap = config['patch_overlap']
if gt_mask_flag == True:
label_path = img_path.replace('image', 'labels')
img_input = cv2.imread(img_path)
if gt_mask_flag == True:
mask_input = cv2.imread(label_path)
if gt_mask_flag == True:
image_patch, mask_patch = patch_gen(img_input, mask_input, p_size,
patch_overlap)
else:
image_patch, mask_patch = patch_gen(img_input, img_input, p_size,
patch_overlap)
val_transform = Compose([
transforms.Resize(config['input_h'], config['input_w']),
transforms.Normalize(),
])
patch_len = len(image_patch)
for idx in range(patch_len):
img = image_patch[idx]
img = cv2.resize(img, (img_size, img_size))
mask = img
if val_transform is not None:
augmented = val_transform(image=img, mask=mask)
img = augmented['image']
img = img.astype('float32') / 255
img = img.transpose(2, 0, 1)
img_patch_set.append(img)
img_patch_set = np.array(img_patch_set)
mask_patch_set = np.array(mask_patch)
return img_input, img_patch_set, mask_patch_set
def get_tta_transforms():
return Compose([
A.RandomResizedCrop(
height=hparams.sz,
width=hparams.sz,
scale=(0.7, 1.0),
),
# AdvancedHairAugmentation(),
A.GridDistortion(),
A.RandomBrightnessContrast(),
A.ShiftScaleRotate(),
A.Flip(p=0.5),
A.CoarseDropout(
max_height=int(hparams.sz / 10),
max_width=int(hparams.sz / 10),
),
# A.HueSaturationValue(),
A.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225],
max_pixel_value=255,
),
ToTensorV2(),
])
img_as_rgb = np.rot90(img_as_rgb)
# Transform image to tensor
if self.transform is not None:
# Apply transformations
augmented = self.transform(image=img_as_rgb)
# Convert numpy array to PIL Image
img_as_tensor = augmented['image']
# Return image and the label
return (img_as_tensor, single_image_label)
#%% with albumentations for Image level classification
albumentations_transform_im = {
'train':
Compose([
ab_transforms.Normalize(mean=[0.5944, 0.4343, 0.5652],
std=[0.2593, 0.2293, 0.2377]),
ToTensor()
]),
'val':
Compose([
ab_transforms.Normalize(mean=[0.5944, 0.4343, 0.5652],
std=[0.2593, 0.2293, 0.2377]),
ToTensor()
]),
}
class Img_CustomDataset(Dataset):
def __init__(self, im_paths, labels, phase=None):
"""
Args:
def main():
args = parse_args()
if args.name is None:
args.name = '%s_%s' % (args.arch, datetime.now().strftime('%m%d%H'))
if not os.path.exists('models/%s' % args.name):
os.makedirs('models/%s' % args.name)
if args.resume:
args = joblib.load('models/%s/args.pkl' % args.name)
args.resume = True
print('Config -----')
for arg in vars(args):
print('- %s: %s' % (arg, getattr(args, arg)))
print('------------')
with open('models/%s/args.txt' % args.name, 'w') as f:
for arg in vars(args):
print('- %s: %s' % (arg, getattr(args, arg)), file=f)
joblib.dump(args, 'models/%s/args.pkl' % args.name)
if args.seed is not None and not args.resume:
print('set random seed')
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if args.loss == 'BCEWithLogitsLoss':
criterion = BCEWithLogitsLoss().cuda()
elif args.loss == 'WeightedBCEWithLogitsLoss':
criterion = BCEWithLogitsLoss(weight=torch.Tensor([1., 1., 1., 1., 1., 2.]),
smooth=args.label_smooth).cuda()
elif args.loss == 'FocalLoss':
criterion = FocalLoss().cuda()
elif args.loss == 'WeightedFocalLoss':
criterion = FocalLoss(weight=torch.Tensor([1., 1., 1., 1., 1., 2.])).cuda()
else:
raise NotImplementedError
if args.pred_type == 'all':
num_outputs = 6
elif args.pred_type == 'except_any':
num_outputs = 5
else:
raise NotImplementedError
cudnn.benchmark = True
# create model
model = get_model(model_name=args.arch,
num_outputs=num_outputs,
freeze_bn=args.freeze_bn,
dropout_p=args.dropout_p,
pooling=args.pooling,
lp_p=args.lp_p)
model = model.cuda()
train_transform = Compose([
transforms.Resize(args.img_size, args.img_size),
transforms.HorizontalFlip() if args.hflip else NoOp(),
transforms.VerticalFlip() if args.vflip else NoOp(),
transforms.ShiftScaleRotate(
shift_limit=args.shift_limit,
scale_limit=args.scale_limit,
rotate_limit=args.rotate_limit,
border_mode=cv2.BORDER_CONSTANT,
value=0,
p=args.shift_scale_rotate_p
) if args.shift_scale_rotate else NoOp(),
transforms.RandomContrast(
limit=args.contrast_limit,
p=args.contrast_p
) if args.contrast else NoOp(),
RandomErase() if args.random_erase else NoOp(),
transforms.CenterCrop(args.crop_size, args.crop_size) if args.center_crop else NoOp(),
ForegroundCenterCrop(args.crop_size) if args.foreground_center_crop else NoOp(),
transforms.RandomCrop(args.crop_size, args.crop_size) if args.random_crop else NoOp(),
transforms.Normalize(mean=model.mean, std=model.std),
ToTensor(),
])
if args.img_type:
stage_1_train_dir = 'processed/stage_1_train_%s' %args.img_type
else:
stage_1_train_dir = 'processed/stage_1_train'
df = pd.read_csv('inputs/stage_1_train.csv')
img_paths = np.array([stage_1_train_dir + '/' + '_'.join(s.split('_')[:-1]) + '.png' for s in df['ID']][::6])
labels = np.array([df.loc[c::6, 'Label'].values for c in range(6)]).T.astype('float32')
df = df[::6]
df['img_path'] = img_paths
for c in range(6):
df['label_%d' %c] = labels[:, c]
df['ID'] = df['ID'].apply(lambda s: '_'.join(s.split('_')[:-1]))
meta_df = pd.read_csv('processed/stage_1_train_meta.csv')
meta_df['ID'] = meta_df['SOPInstanceUID']
test_meta_df = pd.read_csv('processed/stage_1_test_meta.csv')
df = pd.merge(df, meta_df, how='left')
patient_ids = meta_df['PatientID'].unique()
test_patient_ids = test_meta_df['PatientID'].unique()
if args.remove_test_patient_ids:
patient_ids = np.array([s for s in patient_ids if not s in test_patient_ids])
train_img_paths = np.hstack(df[['img_path', 'PatientID']].groupby(['PatientID'])['img_path'].apply(np.array).loc[patient_ids].to_list()).astype('str')
train_labels = []
for c in range(6):
train_labels.append(np.hstack(df[['label_%d' %c, 'PatientID']].groupby(['PatientID'])['label_%d' %c].apply(np.array).loc[patient_ids].to_list()))
train_labels = np.array(train_labels).T
if args.resume:
checkpoint = torch.load('models/%s/checkpoint.pth.tar' % args.name)
# train
train_set = Dataset(
train_img_paths,
train_labels,
transform=train_transform)
train_loader = torch.utils.data.DataLoader(
train_set,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
# pin_memory=True,
)
if args.optimizer == 'Adam':
optimizer = optim.Adam(
filter(lambda p: p.requires_grad, model.parameters()), lr=args.lr, weight_decay=args.weight_decay)
elif args.optimizer == 'AdamW':
optimizer = optim.AdamW(
filter(lambda p: p.requires_grad, model.parameters()), lr=args.lr, weight_decay=args.weight_decay)
elif args.optimizer == 'RAdam':
optimizer = RAdam(
filter(lambda p: p.requires_grad, model.parameters()), lr=args.lr, weight_decay=args.weight_decay)
elif args.optimizer == 'SGD':
optimizer = optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=args.lr,
momentum=args.momentum, weight_decay=args.weight_decay, nesterov=args.nesterov)
else:
raise NotImplementedError
if args.apex:
amp.initialize(model, optimizer, opt_level='O1')
if args.scheduler == 'CosineAnnealingLR':
scheduler = lr_scheduler.CosineAnnealingLR(
optimizer, T_max=args.epochs, eta_min=args.min_lr)
elif args.scheduler == 'MultiStepLR':
scheduler = lr_scheduler.MultiStepLR(optimizer,
milestones=[int(e) for e in args.milestones.split(',')], gamma=args.gamma)
else:
raise NotImplementedError
log = {
'epoch': [],
'loss': [],
}
start_epoch = 0
if args.resume:
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
scheduler.load_state_dict(checkpoint['scheduler'])
start_epoch = checkpoint['epoch']
log = pd.read_csv('models/%s/log.csv' % args.name).to_dict(orient='list')
for epoch in range(start_epoch, args.epochs):
print('Epoch [%d/%d]' % (epoch + 1, args.epochs))
# train for one epoch
train_loss = train(args, train_loader, model, criterion, optimizer, epoch)
if args.scheduler == 'CosineAnnealingLR':
scheduler.step()
print('loss %.4f' % (train_loss))
log['epoch'].append(epoch)
log['loss'].append(train_loss)
pd.DataFrame(log).to_csv('models/%s/log.csv' % args.name, index=False)
torch.save(model.state_dict(), 'models/%s/model.pth' % args.name)
print("=> saved model")
state = {
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
'scheduler': scheduler.state_dict(),
}
torch.save(state, 'models/%s/checkpoint.pth.tar' % args.name)
def AlbumentationTestTransform(self):
tf = tc.Compose([ta.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
tp.ToTensor()
# tp.ToTensor(dict(mean=(0.4914, 0.4822, 0.4465), std=(0.247, 0.2435, 0.2616)))
])
return lambda img: tf(image = np.array(img))["image"]
def get_transforms(phase: str, cli_args) -> Dict[str, Compose]:
"""Get composed albumentations augmentations
Parameters
----------
phase : str
Phase of learning
In ['train', 'val']
cli_args
Arguments coming all the way from `main.py`
Returns
-------
transforms: dict[str, albumentations.core.composition.Compose]
Composed list of transforms
"""
aug_transforms = []
im_sz = (cli_args.image_size, cli_args.image_size)
if phase == "train":
# Data augmentation for training only
aug_transforms.extend([
tf.ShiftScaleRotate(
shift_limit=0,
scale_limit=0.1,
rotate_limit=15,
p=0.5),
tf.Flip(p=0.5),
tf.RandomRotate90(p=0.5),
])
# Exotic Augmentations for train only 🤤
aug_transforms.extend([
tf.RandomBrightnessContrast(p=0.5),
tf.ElasticTransform(p=0.5),
tf.MultiplicativeNoise(multiplier=(0.5, 1.5),
per_channel=True, p=0.2),
])
aug_transforms.extend([
tf.RandomSizedCrop(min_max_height=im_sz,
height=im_sz[0],
width=im_sz[1],
w2h_ratio=1.0,
interpolation=cv2.INTER_LINEAR,
p=1.0),
])
aug_transforms = Compose(aug_transforms)
mask_only_transforms = Compose([
tf.Normalize(mean=0, std=1, always_apply=True)
])
image_only_transforms = Compose([
tf.Normalize(mean=(0.0, 0.0, 0.0), std=(1.0, 1.0, 1.0),
always_apply=True)
])
final_transforms = Compose([
ToTensorV2()
])
transforms = {
'aug': aug_transforms,
'img_only': image_only_transforms,
'mask_only': mask_only_transforms,
'final': final_transforms
}
return transforms
def main():
args = parse_args()
config_file = "../configs/config_SN7.json"
config_dict = json.loads(open(config_file, 'rt').read())
#config_dict = json.loads(open(sys.argv[1], 'rt').read())
file_dict = config_dict['file_path']
val_config = config_dict['val_config']
name = val_config['name']
input_folder =file_dict['input_path'] # '../inputs'
model_folder = file_dict['model_path'] # '../models'
output_folder = file_dict['output_path'] # '../models'
ss_unet_GAN = True
# create model
if ss_unet_GAN == False:
path = os.path.join(model_folder, '%s/config.yml' % name)
with open(os.path.join(model_folder, '%s/config.yml' % name), 'r') as f:
config = yaml.load(f, Loader=yaml.FullLoader)
config['name'] = name
print('-' * 20)
for key in config.keys():
print('%s: %s' % (key, str(config[key])))
print('-' * 20)
cudnn.benchmark = True
print("=> creating model %s" % config['arch'])
model = archs.__dict__[config['arch']](config['num_classes'],
config['input_channels'],
config['deep_supervision'])
model = model.cuda()
#img_ids = glob(os.path.join(input_folder, config['dataset'], 'images', '*' + config['img_ext']))
#img_ids = [os.path.splitext(os.path.basename(p))[0] for p in img_ids]
#_, val_img_ids = train_test_split(img_ids, test_size=0.2, random_state=41)
model_dict = torch.load(os.path.join(model_folder,'%s/model.pth' %config['name']))
if "state_dict" in model_dict.keys():
model_dict = remove_prefix(model_dict['state_dict'], 'module.')
else:
model_dict = remove_prefix(model_dict, 'module.')
model.load_state_dict(model_dict, strict=False)
#model.load_state_dict(torch.load(os.path.join(model_folder,'%s/model.pth' %config['name'])))
model.eval()
else:
val_config = config_dict['val_config']
generator_name = val_config['name']
with open(os.path.join(model_folder, '%s/config.yml' % generator_name), 'r') as f:
config = yaml.load(f, Loader=yaml.FullLoader)
generator = Generator(config)
generator = generator.cuda()
'''
with open(os.path.join(model_folder, '%s/config.yml' % name), 'r') as f:
config = yaml.load(f, Loader=yaml.FullLoader)
'''
config['name'] = name
model_dict = torch.load(os.path.join(model_folder,'%s/model.pth' %config['name']))
if "state_dict" in model_dict.keys():
model_dict = remove_prefix(model_dict['state_dict'], 'module.')
else:
model_dict = remove_prefix(model_dict, 'module.')
generator.load_state_dict(model_dict, strict=False)
#model.load_state_dict(torch.load(os.path.join(model_folder,'%s/model.pth' %config['name'])))
generator.eval()
# Data loading code
img_ids = glob(os.path.join(input_folder, config['val_dataset'], 'images','test', '*' + config['img_ext']))
val_img_ids = [os.path.splitext(os.path.basename(p))[0] for p in img_ids]
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
val_transform = Compose([
transforms.Resize(config['input_h'], config['input_w']),
#transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5)),
transforms.Normalize(mean=mean, std=std),
])
val_dataset = Dataset(
img_ids=val_img_ids,
img_dir=os.path.join(input_folder, config['val_dataset'], 'images','test'),
mask_dir=os.path.join(input_folder, config['val_dataset'], 'annotations','test'),
img_ext=config['img_ext'],
mask_ext=config['mask_ext'],
num_classes=config['num_classes'],
input_channels=config['input_channels'],
transform=val_transform)
val_loader = torch.utils.data.DataLoader(
val_dataset,
batch_size=1, #config['batch_size'],
shuffle=False,
num_workers=config['num_workers'],
drop_last=False)
avg_meters = {'iou': AverageMeter(),
'dice' : AverageMeter()}
num_classes = config['num_classes']
for c in range(config['num_classes']):
os.makedirs(os.path.join( output_folder, config['name'], str(c)), exist_ok=True)
csv_save_name = os.path.join(output_folder, config['name'] + '_result' + '.csv')
result_submission = []
with torch.no_grad():
pbar = tqdm(total=len(val_loader))
for ori_img, input, target, targets, meta in val_loader:
input = input.cuda()
target = target.cuda()
# compute output
if ss_unet_GAN == True:
if config['deep_supervision']:
output = generator(input)[-1]
else:
output = generator(input)
else:
if config['deep_supervision']:
output = model(input)[-1]
else:
output = model(input)
out_m = output[:, 1:num_classes, :, :].clone()
tar_m = target[:, 1:num_classes, :, :].clone()
iou = iou_score(out_m, tar_m)
dice = dice_coef(out_m, tar_m)
result_submission.append([meta['img_id'][0], iou, dice])
avg_meters['iou'].update(iou, input.size(0))
avg_meters['dice'].update(dice, input.size(0))
output = torch.sigmoid(output).cpu().numpy()
masks = target.cpu()
for i in range(len(output)):
for idx_c in range(num_classes):
tmp_mask = np.array(masks[i][idx_c])
mask = np.array(255 * tmp_mask).astype('uint8')
mask_out = np.array(255 * output[i][idx_c]).astype('uint8')
mask_output = np.zeros((mask_out.shape[0], mask_out.shape[1]))
mask_output = mask_output.astype('uint8')
mask_ = mask_out > 127
mask_output[mask_] = 255
if idx_c >0:
save_GT_RE_mask(output_folder, config, meta, idx_c, i, ori_img, mask, mask_output)
postfix = OrderedDict([
('iou', avg_meters['iou'].avg),
('dice', avg_meters['dice'].avg),
])
pbar.set_postfix(postfix)
pbar.update(1)
pbar.close()
result_save_to_csv_filename(csv_save_name, result_submission)
print('IoU: %.4f' % avg_meters['iou'].avg)
print('dice: %.4f' % avg_meters['dice'].avg)
torch.cuda.empty_cache()
def main():
# config = vars(parse_args_func())
#config_file = "../configs/config_v1.json"
args = vars(parse_args_func())
config_file = args['config']
config_dict = json.loads(open(config_file, 'rt').read())
# config_dict = json.loads(open(sys.argv[1], 'rt').read())
file_dict = config_dict['file_path']
config = config_dict['opt_config']
input_folder = file_dict['input_path'] # '../inputs'
checkpoint_folder = file_dict['checkpoint_path'] # '../checkpoint'
model_folder = file_dict['model_path'] # '../models'
if 'False' in config['deep_supervision']:
config['deep_supervision'] = False
else:
config['deep_supervision'] = True
if 'False' in config['nesterov']:
config['nesterov'] = False
else:
config['nesterov'] = True
if 'None' in config['name']:
config['name'] = None
if config['name'] is None:
config['name'] = '%s_%s_segmodel' % (config['dataset'], config['arch'])
os.makedirs(os.path.join(model_folder, '%s' % config['name']),
exist_ok=True)
print('-' * 20)
for key in config:
print('%s: %s' % (key, config[key]))
print('-' * 20)
with open(os.path.join(model_folder, '%s/config.yml' % config['name']),
'w') as f:
yaml.dump(config, f)
# define loss function (criterion)
if config['loss'] == 'BCEWithLogitsLoss':
criterion = nn.BCEWithLogitsLoss().cuda()
else:
criterion = losses.__dict__[config['loss']]().cuda()
cudnn.benchmark = True
# create model
if 'False' in config['resume']:
config['resume'] = False
else:
config['resume'] = True
# Data loading code
img_ids = glob(
os.path.join(input_folder, config['dataset'], 'images', 'training',
'*' + config['img_ext']))
train_img_ids = [os.path.splitext(os.path.basename(p))[0] for p in img_ids]
img_ids = glob(
os.path.join(input_folder, config['val_dataset'], 'images',
'validation', '*' + config['img_ext']))
val_img_ids = [os.path.splitext(os.path.basename(p))[0] for p in img_ids]
img_ids = glob(
os.path.join(input_folder, config['val_dataset'], 'images', 'test',
'*' + config['img_ext']))
test_img_ids = [os.path.splitext(os.path.basename(p))[0] for p in img_ids]
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
train_transform = Compose([
# transforms.RandomScale ([config['scale_min'], config['scale_max']]),
# transforms.RandomRotate90(),
transforms.Rotate([config['rotate_min'], config['rotate_max']],
value=mean,
mask_value=0),
# transforms.GaussianBlur (),
transforms.Flip(),
# transforms.HorizontalFlip (),
transforms.HueSaturationValue(hue_shift_limit=10,
sat_shift_limit=10,
val_shift_limit=10),
transforms.RandomBrightnessContrast(brightness_limit=0.10,
contrast_limit=0.10,
brightness_by_max=True),
transforms.Resize(config['input_h'], config['input_w']),
transforms.Normalize(mean=mean, std=std),
])
val_transform = Compose([
transforms.Resize(config['input_h'], config['input_w']),
transforms.Normalize(mean=mean, std=std),
])
train_dataset = Dataset(img_ids=train_img_ids,
img_dir=os.path.join(input_folder,
config['dataset'], 'images',
'training'),
mask_dir=os.path.join(input_folder,
config['dataset'],
'annotations', 'training'),
img_ext=config['img_ext'],
mask_ext=config['mask_ext'],
num_classes=config['num_classes'],
input_channels=config['input_channels'],
transform=train_transform)
val_dataset = Dataset(img_ids=val_img_ids,
img_dir=os.path.join(input_folder, config['dataset'],
'images', 'validation'),
mask_dir=os.path.join(input_folder,
config['dataset'],
'annotations', 'validation'),
img_ext=config['img_ext'],
mask_ext=config['mask_ext'],
num_classes=config['num_classes'],
input_channels=config['input_channels'],
transform=val_transform)
test_dataset = Dataset(img_ids=val_img_ids,
img_dir=os.path.join(input_folder,
config['dataset'], 'images',
'test'),
mask_dir=os.path.join(input_folder,
config['dataset'],
'annotations', 'test'),
img_ext=config['img_ext'],
mask_ext=config['mask_ext'],
num_classes=config['num_classes'],
input_channels=config['input_channels'],
transform=val_transform)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=config['batch_size'],
shuffle=True,
num_workers=config['num_workers'],
drop_last=True)
val_loader = torch.utils.data.DataLoader(
val_dataset,
batch_size=1, # config['batch_size'],
shuffle=False,
num_workers=config['num_workers'],
drop_last=False)
test_loader = torch.utils.data.DataLoader(
test_dataset,
batch_size=1, # config['batch_size'],
shuffle=False,
num_workers=config['num_workers'],
drop_last=False)
log = OrderedDict([
('epoch', []),
('lr', []),
('loss', []),
('iou', []),
('dice', []),
('val_loss', []),
('val_iou', []),
('val_dice', []),
])
if not os.path.isdir(checkpoint_folder):
os.mkdir(checkpoint_folder)
# create generator model
#val_config = config_dict['config']
generator_name = config['generator_name']
with open(os.path.join(model_folder, '%s/config.yml' % generator_name),
'r') as f:
g_config = yaml.load(f, Loader=yaml.FullLoader)
generator = Generator(g_config)
generator.initialize_with_srresnet(model_folder, g_config)
lr = config['gan_lr']
# Initialize generator's optimizer
optimizer_g = torch.optim.Adam(params=filter(lambda p: p.requires_grad,
generator.parameters()),
lr=lr)
#params = filter(lambda p: p.requires_grad, generator.parameters())
#optimizer_g, scheduler_g = optimizer_scheduler(params, config)
# Discriminator
# Discriminator parameters
num_classes = config['num_classes']
kernel_size_d = 3 # kernel size in all convolutional blocks
n_channels_d = 64 # number of output channels in the first convolutional block, after which it is doubled in every 2nd block thereafter
n_blocks_d = 8 # number of convolutional blocks
fc_size_d = 1024 # size of the first fully connected layer
discriminator = Discriminator(num_classes,
kernel_size=kernel_size_d,
n_channels=n_channels_d,
n_blocks=n_blocks_d,
fc_size=fc_size_d)
# Initialize discriminator's optimizer
optimizer_d = torch.optim.Adam(params=filter(lambda p: p.requires_grad,
discriminator.parameters()),
lr=lr)
#params = filter(lambda p: p.requires_grad, discriminator.parameters())
#optimizer_d, scheduler_d = optimizer_scheduler(params, config)
adversarial_loss_criterion = nn.BCEWithLogitsLoss()
content_loss_criterion = nn.MSELoss()
generator = generator.cuda()
discriminator = discriminator.cuda()
#truncated_vgg19 = truncated_vgg19.cuda()
content_loss_criterion = content_loss_criterion.cuda()
adversarial_loss_criterion = adversarial_loss_criterion.cuda()
generator = torch.nn.DataParallel(generator)
discriminator = torch.nn.DataParallel(discriminator)
if not os.path.isdir(checkpoint_folder):
os.mkdir(checkpoint_folder)
log_name = config['name']
log_dir = os.path.join(checkpoint_folder, log_name)
writer = SummaryWriter(logdir=log_dir)
best_iou = 0
trigger = 0
Best_dice = 0
iou_AtBestDice = 0
start_epoch = 0
for epoch in range(start_epoch, config['epochs']):
print('Epoch [%d/%d]' % (epoch, config['epochs']))
# train for one epoch
train_log = train(epoch, config, train_loader, generator,
discriminator, criterion, adversarial_loss_criterion,
content_loss_criterion, optimizer_g, optimizer_d)
# evaluate on validation set
val_log = validate(config, val_loader, generator, criterion)
test_log = validate(config, test_loader, generator, criterion)
if Best_dice < test_log['dice']:
Best_dice = test_log['dice']
iou_AtBestDice = test_log['iou']
print(
'loss %.4f - iou %.4f - dice %.4f - val_loss %.4f - val_iou %.4f - val_dice %.4f - test_iou %.4f - test_dice %.4f - Best_dice %.4f - iou_AtBestDice %.4f'
% (train_log['loss'], train_log['iou'], train_log['dice'],
val_log['loss'], val_log['iou'], val_log['dice'],
test_log['iou'], test_log['dice'], Best_dice, iou_AtBestDice))
save_tensorboard(writer, train_log, val_log, test_log, epoch)
log['epoch'].append(epoch)
log['lr'].append(config['lr'])
log['loss'].append(train_log['loss'])
log['iou'].append(train_log['iou'])
log['dice'].append(train_log['dice'])
log['val_loss'].append(val_log['loss'])
log['val_iou'].append(val_log['iou'])
log['val_dice'].append(val_log['dice'])
pd.DataFrame(log).to_csv(os.path.join(model_folder,
'%s/log.csv' % config['name']),
index=False)
trigger += 1
if test_log['iou'] > best_iou:
torch.save(
generator.state_dict(),
os.path.join(model_folder, '%s/model.pth' % config['name']))
best_iou = test_log['iou']
print("=> saved best model")
trigger = 0
# early stopping
if config['early_stopping'] >= 0 and trigger >= config[
'early_stopping']:
print("=> early stopping")
break
torch.cuda.empty_cache()
def main():
config = vars(parse_args())
if config['name'] is None:
if config['deep_supervision']:
config['name'] = '%s_%s_wDS' % (config['dataset'], config['arch'])
else:
config['name'] = '%s_%s_woDS' % (config['dataset'], config['arch'])
os.makedirs('models/%s' % config['name'], exist_ok=True)
print('-' * 20)
for key in config:
print('%s: %s' % (key, config[key]))
print('-' * 20)
with open('models/%s/config.yml' % config['name'], 'w') as f:
yaml.dump(config, f)
# define loss function (criterion)
if config['loss'] == 'BCEWithLogitsLoss':
criterion = nn.BCEWithLogitsLoss().cuda()
else:
criterion = losses.__dict__[config['loss']]().cuda()
cudnn.benchmark = True
# create model
print("=> creating model %s" % config['arch'])
model = archs.__dict__[config['arch']](config['num_classes'],
config['input_channels'],
config['deep_supervision'])
model = model.cuda()
params = filter(lambda p: p.requires_grad, model.parameters())
if config['optimizer'] == 'Adam':
optimizer = optim.Adam(params,
lr=config['lr'],
weight_decay=config['weight_decay'])
elif config['optimizer'] == 'SGD':
optimizer = optim.SGD(params,
lr=config['lr'],
momentum=config['momentum'],
nesterov=config['nesterov'],
weight_decay=config['weight_decay'])
else:
raise NotImplementedError
if config['scheduler'] == 'CosineAnnealingLR':
scheduler = lr_scheduler.CosineAnnealingLR(optimizer,
T_max=config['epochs'],
eta_min=config['min_lr'])
elif config['scheduler'] == 'ReduceLROnPlateau':
scheduler = lr_scheduler.ReduceLROnPlateau(optimizer,
factor=config['factor'],
patience=config['patience'],
verbose=1,
min_lr=config['min_lr'])
elif config['scheduler'] == 'MultiStepLR':
scheduler = lr_scheduler.MultiStepLR(
optimizer,
milestones=[int(e) for e in config['milestones'].split(',')],
gamma=config['gamma'])
elif config['scheduler'] == 'ConstantLR':
scheduler = None
else:
raise NotImplementedError
# Data loading code
img_ids = glob(
os.path.join('inputs', config['dataset'], 'images',
'*' + config['img_ext']))
img_ids = [os.path.splitext(os.path.basename(p))[0] for p in img_ids]
train_img_ids, val_img_ids = train_test_split(img_ids,
test_size=0.2,
random_state=41)
#######################
train_img_ids = img_ids
#######################
train_transform = Compose([
transforms.RandomRotate90(),
transforms.Flip(),
OneOf([
transforms.HueSaturationValue(),
transforms.RandomBrightness(),
transforms.RandomContrast(),
],
p=1),
transforms.Resize(config['input_h'], config['input_w']),
transforms.Normalize(),
])
val_transform = Compose([
transforms.Resize(config['input_h'], config['input_w']),
transforms.Normalize(),
])
train_dataset = Dataset(img_ids=train_img_ids,
img_dir=os.path.join('inputs', config['dataset'],
'images'),
mask_dir=os.path.join('inputs', config['dataset'],
'masks'),
img_ext=config['img_ext'],
mask_ext=config['mask_ext'],
num_classes=config['num_classes'],
transform=train_transform)
val_dataset = Dataset(img_ids=val_img_ids,
img_dir=os.path.join('inputs', config['dataset'],
'images'),
mask_dir=os.path.join('inputs', config['dataset'],
'masks'),
img_ext=config['img_ext'],
mask_ext=config['mask_ext'],
num_classes=config['num_classes'],
transform=val_transform)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=config['batch_size'],
shuffle=True,
num_workers=config['num_workers'],
drop_last=True)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=config['batch_size'],
shuffle=False,
num_workers=config['num_workers'],
drop_last=False)
log = OrderedDict([
('epoch', []),
('lr', []),
('loss', []),
('iou', []),
('val_loss', []),
('val_iou', []),
])
best_iou = 0
trigger = 0
for epoch in range(config['epochs']):
print('Epoch [%d/%d]' % (epoch, config['epochs']))
# train for one epoch
train_log = train(config, train_loader, model, criterion, optimizer)
# evaluate on validation set
val_log = validate(config, val_loader, model, criterion)
if config['scheduler'] == 'CosineAnnealingLR':
scheduler.step()
elif config['scheduler'] == 'ReduceLROnPlateau':
scheduler.step(val_log['loss'])
print('loss %.4f - iou %.4f - val_loss %.4f - val_iou %.4f' %
(train_log['loss'], train_log['iou'], val_log['loss'],
val_log['iou']))
log['epoch'].append(epoch)
log['lr'].append(config['lr'])
log['loss'].append(train_log['loss'])
log['iou'].append(train_log['iou'])
log['val_loss'].append(val_log['loss'])
log['val_iou'].append(val_log['iou'])
pd.DataFrame(log).to_csv('models/%s/log.csv' % config['name'],
index=False)
trigger += 1
if val_log['iou'] > best_iou:
torch.save(model.state_dict(),
'models/%s/model.pth' % config['name'])
best_iou = val_log['iou']
print("=> saved best model")
trigger = 0
# early stopping
if config['early_stopping'] >= 0 and trigger >= config[
'early_stopping']:
print("=> early stopping")
break
torch.cuda.empty_cache()
def main():
config = vars(parse_args())
if config['name'] is None:
config['name'] = '%s_%s' % (config['arch'], datetime.now().strftime('%m%d%H'))
if not os.path.exists('models/pose/%s' % config['name']):
os.makedirs('models/pose/%s' % config['name'])
if config['resume']:
with open('models/pose/%s/config.yml' % config['name'], 'r') as f:
config = yaml.load(f, Loader=yaml.FullLoader)
config['resume'] = True
with open('models/pose/%s/config.yml' % config['name'], 'w') as f:
yaml.dump(config, f)
print('-'*20)
for key in config.keys():
print('- %s: %s' % (key, str(config[key])))
print('-'*20)
cudnn.benchmark = True
df = pd.read_csv('inputs/train.csv')
img_ids = df['ImageId'].values
pose_df = pd.read_csv('processed/pose_train.csv')
pose_df['img_path'] = 'processed/pose_images/train/' + pose_df['img_path']
if config['resume']:
checkpoint = torch.load('models/pose/%s/checkpoint.pth.tar' % config['name'])
if config['rot'] == 'eular':
num_outputs = 3
elif config['rot'] == 'trig':
num_outputs = 6
elif config['rot'] == 'quat':
num_outputs = 4
else:
raise NotImplementedError
if config['loss'] == 'L1Loss':
criterion = nn.L1Loss().cuda()
elif config['loss'] == 'MSELoss':
criterion = nn.MSELoss().cuda()
else:
raise NotImplementedError
train_transform = Compose([
transforms.ShiftScaleRotate(
shift_limit=config['shift_limit'],
scale_limit=0,
rotate_limit=0,
border_mode=cv2.BORDER_CONSTANT,
value=0,
p=config['shift_p']
) if config['shift'] else NoOp(),
OneOf([
transforms.HueSaturationValue(
hue_shift_limit=config['hue_limit'],
sat_shift_limit=config['sat_limit'],
val_shift_limit=config['val_limit'],
p=config['hsv_p']
) if config['hsv'] else NoOp(),
transforms.RandomBrightness(
limit=config['brightness_limit'],
p=config['brightness_p'],
) if config['brightness'] else NoOp(),
transforms.RandomContrast(
limit=config['contrast_limit'],
p=config['contrast_p'],
) if config['contrast'] else NoOp(),
], p=1),
transforms.ISONoise(
p=config['iso_noise_p'],
) if config['iso_noise'] else NoOp(),
transforms.CLAHE(
p=config['clahe_p'],
) if config['clahe'] else NoOp(),
transforms.Resize(config['input_w'], config['input_h']),
transforms.Normalize(),
ToTensor(),
])
val_transform = Compose([
transforms.Resize(config['input_w'], config['input_h']),
transforms.Normalize(),
ToTensor(),
])
folds = []
best_losses = []
kf = KFold(n_splits=config['n_splits'], shuffle=True, random_state=41)
for fold, (train_idx, val_idx) in enumerate(kf.split(img_ids)):
print('Fold [%d/%d]' %(fold + 1, config['n_splits']))
if (config['resume'] and fold < checkpoint['fold'] - 1) or (not config['resume'] and os.path.exists('pose_models/%s/model_%d.pth' % (config['name'], fold+1))):
log = pd.read_csv('models/pose/%s/log_%d.csv' %(config['name'], fold+1))
best_loss = log.loc[log['val_loss'].values.argmin(), 'val_loss']
# best_loss, best_score = log.loc[log['val_loss'].values.argmin(), ['val_loss', 'val_score']].values
folds.append(str(fold + 1))
best_losses.append(best_loss)
# best_scores.append(best_score)
continue
train_img_ids, val_img_ids = img_ids[train_idx], img_ids[val_idx]
train_img_paths = []
train_labels = []
for img_id in train_img_ids:
tmp = pose_df.loc[pose_df.ImageId == img_id]
img_path = tmp['img_path'].values
train_img_paths.append(img_path)
yaw = tmp['yaw'].values
pitch = tmp['pitch'].values
roll = tmp['roll'].values
roll = rotate(roll, np.pi)
if config['rot'] == 'eular':
label = np.array([
yaw,
pitch,
roll
]).T
elif config['rot'] == 'trig':
label = np.array([
np.cos(yaw),
np.sin(yaw),
np.cos(pitch),
np.sin(pitch),
np.cos(roll),
np.sin(roll),
]).T
elif config['rot'] == 'quat':
raise NotImplementedError
else:
raise NotImplementedError
train_labels.append(label)
train_img_paths = np.hstack(train_img_paths)
train_labels = np.vstack(train_labels)
val_img_paths = []
val_labels = []
for img_id in val_img_ids:
tmp = pose_df.loc[pose_df.ImageId == img_id]
img_path = tmp['img_path'].values
val_img_paths.append(img_path)
yaw = tmp['yaw'].values
pitch = tmp['pitch'].values
roll = tmp['roll'].values
roll = rotate(roll, np.pi)
if config['rot'] == 'eular':
label = np.array([
yaw,
pitch,
roll
]).T
elif config['rot'] == 'trig':
label = np.array([
np.cos(yaw),
np.sin(yaw),
np.cos(pitch),
np.sin(pitch),
np.cos(roll),
np.sin(roll),
]).T
elif config['rot'] == 'quat':
raise NotImplementedError
else:
raise NotImplementedError
val_labels.append(label)
val_img_paths = np.hstack(val_img_paths)
val_labels = np.vstack(val_labels)
# train
train_set = PoseDataset(
train_img_paths,
train_labels,
transform=train_transform,
)
train_loader = torch.utils.data.DataLoader(
train_set,
batch_size=config['batch_size'],
shuffle=True,
num_workers=config['num_workers'],
# pin_memory=True,
)
val_set = PoseDataset(
val_img_paths,
val_labels,
transform=val_transform,
)
val_loader = torch.utils.data.DataLoader(
val_set,
batch_size=config['batch_size'],
shuffle=False,
num_workers=config['num_workers'],
# pin_memory=True,
)
# create model
model = get_pose_model(config['arch'],
num_outputs=num_outputs,
freeze_bn=config['freeze_bn'])
model = model.cuda()
params = filter(lambda p: p.requires_grad, model.parameters())
if config['optimizer'] == 'Adam':
optimizer = optim.Adam(params, lr=config['lr'], weight_decay=config['weight_decay'])
elif config['optimizer'] == 'AdamW':
optimizer = optim.AdamW(params, lr=config['lr'], weight_decay=config['weight_decay'])
elif config['optimizer'] == 'RAdam':
optimizer = RAdam(params, lr=config['lr'], weight_decay=config['weight_decay'])
elif config['optimizer'] == 'SGD':
optimizer = optim.SGD(params, lr=config['lr'], momentum=config['momentum'],
nesterov=config['nesterov'], weight_decay=config['weight_decay'])
else:
raise NotImplementedError
if config['scheduler'] == 'CosineAnnealingLR':
scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=config['epochs'], eta_min=config['min_lr'])
elif config['scheduler'] == 'ReduceLROnPlateau':
scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, factor=config['factor'], patience=config['patience'],
verbose=1, min_lr=config['min_lr'])
elif config['scheduler'] == 'MultiStepLR':
scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=[int(e) for e in config['milestones'].split(',')], gamma=config['gamma'])
else:
raise NotImplementedError
log = {
'epoch': [],
'loss': [],
# 'score': [],
'val_loss': [],
# 'val_score': [],
}
best_loss = float('inf')
# best_score = float('inf')
start_epoch = 0
if config['resume'] and fold == checkpoint['fold'] - 1:
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
scheduler.load_state_dict(checkpoint['scheduler'])
start_epoch = checkpoint['epoch']
log = pd.read_csv('models/pose/%s/log_%d.csv' % (config['name'], fold+1)).to_dict(orient='list')
best_loss = checkpoint['best_loss']
for epoch in range(start_epoch, config['epochs']):
print('Epoch [%d/%d]' % (epoch + 1, config['epochs']))
# train for one epoch
train_loss = train(config, train_loader, model, criterion, optimizer, epoch)
# evaluate on validation set
val_loss = validate(config, val_loader, model, criterion)
if config['scheduler'] == 'CosineAnnealingLR':
scheduler.step()
elif config['scheduler'] == 'ReduceLROnPlateau':
scheduler.step(val_loss)
print('loss %.4f - val_loss %.4f' % (train_loss, val_loss))
# print('loss %.4f - score %.4f - val_loss %.4f - val_score %.4f'
# % (train_loss, train_score, val_loss, val_score))
log['epoch'].append(epoch)
log['loss'].append(train_loss)
# log['score'].append(train_score)
log['val_loss'].append(val_loss)
# log['val_score'].append(val_score)
pd.DataFrame(log).to_csv('models/pose/%s/log_%d.csv' % (config['name'], fold+1), index=False)
if val_loss < best_loss:
torch.save(model.state_dict(), 'models/pose/%s/model_%d.pth' % (config['name'], fold+1))
best_loss = val_loss
# best_score = val_score
print("=> saved best model")
state = {
'fold': fold + 1,
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'best_loss': best_loss,
'optimizer': optimizer.state_dict(),
'scheduler': scheduler.state_dict(),
}
torch.save(state, 'models/pose/%s/checkpoint.pth.tar' % config['name'])
print('val_loss: %f' % best_loss)
# print('val_score: %f' % best_score)
folds.append(str(fold + 1))
best_losses.append(best_loss)
# best_scores.append(best_score)
results = pd.DataFrame({
'fold': folds + ['mean'],
'best_loss': best_losses + [np.mean(best_losses)],
# 'best_score': best_scores + [np.mean(best_scores)],
})
print(results)
results.to_csv('models/pose/%s/results.csv' % config['name'], index=False)
del model
torch.cuda.empty_cache()
del train_set, train_loader
del val_set, val_loader
gc.collect()
if not config['cv']:
break
def main():
args = parse_args()
with open('models/%s/config.yml' % args.name, 'r') as f:
config = yaml.load(f, Loader=yaml.FullLoader)
print('-' * 20)
for key in config.keys():
print('%s: %s' % (key, str(config[key])))
print('-' * 20)
cudnn.benchmark = True
# create model
print("=> creating model %s" % config['arch'])
model = archs.__dict__[config['arch']](config['num_classes'],
config['input_channels'],
config['deep_supervision'])
model = model.cuda()
# Data loading code
img_ids = glob(
os.path.join('inputs', config['dataset'], 'test\\images',
'*' + config['img_ext'])) ############ 바꿈
img_ids = [os.path.splitext(os.path.basename(p))[0] for p in img_ids]
val_img_ids = img_ids
model.load_state_dict(torch.load('models/%s/model.pth' % config['name']))
model.eval()
val_transform = Compose([
transforms.Resize(config['input_h'], config['input_w']),
transforms.Normalize(),
])
val_dataset = Dataset(
img_ids=val_img_ids,
img_dir=os.path.join('inputs', config['dataset'],
'test\\images'), ############ 바꿈
mask_dir=os.path.join('inputs', config['dataset'],
'test\\masks'), ############ 바꿈
img_ext=config['img_ext'],
mask_ext=config['mask_ext'],
num_classes=config['num_classes'],
transform=val_transform)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=config['batch_size'],
shuffle=False,
num_workers=config['num_workers'],
drop_last=False)
avg_meter = AverageMeter()
for c in range(config['num_classes']):
os.makedirs(os.path.join('outputs', config['name'], str(c)),
exist_ok=True)
with torch.no_grad():
for input, target, meta in tqdm(val_loader, total=len(val_loader)):
input = input.cuda()
target = target.cuda()
# compute output
if config['deep_supervision']:
output = model(input)[-1]
else:
output = model(input)
iou = iou_score(output, target)
avg_meter.update(iou, input.size(0))
output = torch.sigmoid(output).cpu().numpy()
for i in range(len(output)):
for c in range(config['num_classes']):
cv2.imwrite(
os.path.join('outputs', config['name'], str(c),
meta['img_id'][i] + '.jpg'),
(output[i, c] * 255).astype('uint8'))
print('IoU: %.4f' % avg_meter.avg)
torch.cuda.empty_cache()
def Normalize(mean=(0.5, 0.5, 0.5),
std=(0.5, 0.5, 0.5), max_pixel_value=255.0):
return transforms.Normalize(mean=mean,
std=std,
max_pixel_value=max_pixel_value)
def main():
args = parse_args()
with open('models/pose/%s/config.yml' % args.pose_name, 'r') as f:
config = yaml.load(f, Loader=yaml.FullLoader)
print('-' * 20)
for key in config.keys():
print('%s: %s' % (key, str(config[key])))
print('-' * 20)
cudnn.benchmark = True
df = pd.read_csv('inputs/train.csv')
img_ids = df['ImageId'].values
img_paths = np.array('inputs/train_images/' + df['ImageId'].values +
'.jpg')
mask_paths = np.array('inputs/train_masks/' + df['ImageId'].values +
'.jpg')
labels = np.array(
[convert_str_to_labels(s) for s in df['PredictionString']])
with open('outputs/decoded/val/%s.json' % args.det_name, 'r') as f:
dets = json.load(f)
if config['rot'] == 'eular':
num_outputs = 3
elif config['rot'] == 'trig':
num_outputs = 6
elif config['rot'] == 'quat':
num_outputs = 4
else:
raise NotImplementedError
test_transform = Compose([
transforms.Resize(config['input_w'], config['input_h']),
transforms.Normalize(),
ToTensor(),
])
det_df = {
'ImageId': [],
'img_path': [],
'det': [],
'mask': [],
}
name = '%s_%.2f' % (args.det_name, args.score_th)
if args.nms:
name += '_nms%.2f' % args.nms_th
output_dir = 'processed/pose_images/val/%s' % name
os.makedirs(output_dir, exist_ok=True)
df = []
kf = KFold(n_splits=config['n_splits'], shuffle=True, random_state=41)
for fold, (train_idx, val_idx) in enumerate(kf.split(img_paths)):
print('Fold [%d/%d]' % (fold + 1, config['n_splits']))
# create model
model = get_pose_model(config['arch'],
num_outputs=num_outputs,
freeze_bn=config['freeze_bn'])
model = model.cuda()
model_path = 'models/pose/%s/model_%d.pth' % (config['name'], fold + 1)
if not os.path.exists(model_path):
print('%s is not exists.' % model_path)
continue
model.load_state_dict(torch.load(model_path))
model.eval()
val_img_ids = img_ids[val_idx]
val_img_paths = img_paths[val_idx]
fold_det_df = {
'ImageId': [],
'img_path': [],
'det': [],
'mask': [],
}
for img_id, img_path in tqdm(zip(val_img_ids, val_img_paths),
total=len(val_img_ids)):
img = cv2.imread(img_path)
height, width = img.shape[:2]
det = np.array(dets[img_id])
det = det[det[:, 6] > args.score_th]
if args.nms:
det = nms(det, dist_th=args.nms_th)
for k in range(len(det)):
pitch, yaw, roll, x, y, z, score, w, h = det[k]
fold_det_df['ImageId'].append(img_id)
fold_det_df['det'].append(det[k])
output_path = '%s_%d.jpg' % (img_id, k)
fold_det_df['img_path'].append(output_path)
x, y = convert_3d_to_2d(x, y, z)
w *= 1.1
h *= 1.1
xmin = int(round(x - w / 2))
xmax = int(round(x + w / 2))
ymin = int(round(y - h / 2))
ymax = int(round(y + h / 2))
cropped_img = img[ymin:ymax, xmin:xmax]
if cropped_img.shape[0] > 0 and cropped_img.shape[1] > 0:
cv2.imwrite(os.path.join(output_dir, output_path),
cropped_img)
fold_det_df['mask'].append(1)
else:
fold_det_df['mask'].append(0)
fold_det_df = pd.DataFrame(fold_det_df)
test_set = PoseDataset(output_dir + '/' +
fold_det_df['img_path'].values,
fold_det_df['det'].values,
transform=test_transform,
masks=fold_det_df['mask'].values)
test_loader = torch.utils.data.DataLoader(
test_set,
batch_size=config['batch_size'],
shuffle=False,
num_workers=config['num_workers'],
# pin_memory=True,
)
fold_dets = []
with torch.no_grad():
for input, batch_det, mask in tqdm(test_loader,
total=len(test_loader)):
input = input.cuda()
batch_det = batch_det.numpy()
mask = mask.numpy()
output = model(input)
output = output.cpu()
if config['rot'] == 'trig':
yaw = torch.atan2(output[..., 1:2], output[..., 0:1])
pitch = torch.atan2(output[..., 3:4], output[..., 2:3])
roll = torch.atan2(output[..., 5:6], output[..., 4:5])
roll = rotate(roll, -np.pi)
pitch = pitch.cpu().numpy()[:, 0]
yaw = yaw.cpu().numpy()[:, 0]
roll = roll.cpu().numpy()[:, 0]
batch_det[mask, 0] = pitch[mask]
batch_det[mask, 1] = yaw[mask]
batch_det[mask, 2] = roll[mask]
fold_dets.append(batch_det)
fold_dets = np.vstack(fold_dets)
fold_det_df['det'] = fold_dets.tolist()
fold_det_df = fold_det_df.groupby('ImageId')['det'].apply(list)
fold_det_df = pd.DataFrame({
'ImageId': fold_det_df.index.values,
'PredictionString': fold_det_df.values,
})
df.append(fold_det_df)
break
df = pd.concat(df).reset_index(drop=True)
for i in tqdm(range(len(df))):
img_id = df.loc[i, 'ImageId']
det = np.array(df.loc[i, 'PredictionString'])
if args.show:
img = cv2.imread('inputs/train_images/%s.jpg' % img_id)
img_pred = visualize(img, det)
plt.imshow(img_pred[..., ::-1])
plt.show()
df.loc[i, 'PredictionString'] = convert_labels_to_str(det[:, :7])
name += '_%s' % args.pose_name
df.to_csv('outputs/submissions/val/%s.csv' % name, index=False)