def __init__(self,
df=None,
size=150,
size_gs=224,
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225],
augment=True,
frames=15,
wlen=0,
stochastic=True):
""" Dataset initialization
Parameters
----------
df : pd.DataFrame
Dataframe with preprocessed data
transform : torchvision.transforms
Transformation operations for loaded images
path : str
Path to folder with the data
frames : int
Frames to load per video
"""
assert df is not None, 'Missing dataframe for data'
self.frames = frames
self.df = df[(df['nframes'] >= frames) & (df['wlen'] > wlen)]
self.wlen = wlen
self.stochastic = stochastic
addtl_img = {}
for idx in range(frames):
addtl_img['image{}'.format(idx)] = 'image'
if augment:
self.transform = albumentations.Compose(
[
ShiftScaleRotate(p=0.3,
scale_limit=0.25,
border_mode=1,
rotate_limit=15),
HorizontalFlip(p=0.2),
RandomBrightnessContrast(
p=0.3, brightness_limit=0.25, contrast_limit=0.5),
MotionBlur(p=.2),
GaussNoise(p=.2),
JpegCompression(p=.2, quality_lower=50),
Normalize(mean, std)
],
additional_targets=addtl_img)
self.transform_gs = albumentations.Compose(
[
ShiftScaleRotate(p=0.3,
scale_limit=0.25,
border_mode=1,
rotate_limit=15),
HorizontalFlip(p=0.2),
RandomBrightnessContrast(
p=0.3, brightness_limit=0.25, contrast_limit=0.5),
MotionBlur(p=.2),
GaussNoise(p=.2),
JpegCompression(p=.2, quality_lower=50)
],
additional_targets=addtl_img)
else:
self.transform = albumentations.Compose([Normalize(mean, std)])
self.transform_gs = None
self.resize = transforms.Resize((size, size))
self.resize_gs = transforms.Compose(
[transforms.Resize((size_gs, size_gs)),
transforms.Grayscale()])
X, val_X, y, val_y = get_random_sampling(paths, y, val_paths, val_y)
print('There are ' + str(y.count(1)) + ' fake train samples')
print('There are ' + str(y.count(0)) + ' real train samples')
print('There are ' + str(val_y.count(1)) + ' fake val samples')
print('There are ' + str(val_y.count(0)) + ' real val samples')
import albumentations
from albumentations.augmentations.transforms import ShiftScaleRotate, HorizontalFlip, Normalize, RandomBrightnessContrast, MotionBlur, Blur, GaussNoise, JpegCompression, CoarseDropout
train_transform = albumentations.Compose([
ShiftScaleRotate(p=0.3, scale_limit=0.25, border_mode=1, rotate_limit=25),
HorizontalFlip(p=0.2),
RandomBrightnessContrast(p=0.3, brightness_limit=0.25, contrast_limit=0.5),
MotionBlur(p=.2),
GaussNoise(p=.2),
JpegCompression(p=.2, quality_lower=50),
# CoarseDropout(max_holes=8, max_height=8, max_width=8, fill_value=0, always_apply=False, p=0.5),
Normalize()
])
val_transform = albumentations.Compose([Normalize()])
train_dataset = ImageDataset(X, y, transform=train_transform)
val_dataset = ImageDataset(val_X, val_y, transform=val_transform)
### Train
import gc
history = pd.DataFrame()
history2 = pd.DataFrame()
def main():
args = arg_parser()
seed_everything(args.seed)
if cuda.is_available() and not args.cpu:
device = torch.device("cuda:0")
else:
device = torch.device("cpu")
print(device)
if args.model_type == 'cnn':
if args.preprocess:
train_df = pd.read_csv('../input/preprocessed_train_df.csv')
valid_df = pd.read_csv('../input/preprocessed_valid_df.csv')
else:
train_df = pd.read_csv('../input/train_df.csv')
valid_df = pd.read_csv('../input/valid_df.csv')
valid_sample_num = 40000
elif args.model_type == 'lrcn':
if args.preprocess:
train_df = pd.read_pickle(
'../input/preprocessed_lrcn_train_df.pkl')
valid_df = pd.read_pickle(
'../input/preprocessed_lrcn_train_df.pkl')
else:
train_df = pd.read_pickle('../input/lrcn_train_df.pkl')
valid_df = pd.read_pickle('../input/lrcn_valid_df.pkl')
valid_sample_num = 15000
print("number of train data {}".format(len(train_df)))
print("number of valid data {}\n".format(len(valid_df)))
train_df = train_df.sample(frac=args.train_sample_num,
random_state=args.seed).reset_index(drop=True)
valid_df_sub = valid_df.sample(
frac=1.0, random_state=42).reset_index(drop=True)[:valid_sample_num]
valid_df_sub1 = valid_df.sample(
frac=1.0, random_state=52).reset_index(drop=True)[:valid_sample_num]
valid_df_sub2 = valid_df.sample(
frac=1.0, random_state=62).reset_index(drop=True)[:valid_sample_num]
del valid_df
gc.collect()
if args.DEBUG:
train_df = train_df[:1000]
valid_df_sub = valid_df_sub[:1000]
valid_df_sub1 = valid_df_sub1[:1000]
valid_df_sub2 = valid_df_sub2[:1000]
if args.model_type == 'cnn':
train_transforms = albumentations.Compose([
HorizontalFlip(p=0.3),
# ShiftScaleRotate(p=0.3, scale_limit=0.25, border_mode=1, rotate_limit=25),
# RandomBrightnessContrast(p=0.2, brightness_limit=0.25, contrast_limit=0.5),
# MotionBlur(p=0.2),
GaussNoise(p=0.3),
JpegCompression(p=0.3, quality_lower=50),
# Normalize()
])
valid_transforms = albumentations.Compose([
HorizontalFlip(p=0.2),
albumentations.OneOf([
JpegCompression(quality_lower=8, quality_upper=30, p=1.0),
GaussNoise(p=1.0),
],
p=0.22),
# Normalize()
])
elif args.model_type == 'lrcn':
train_transforms = None
valid_transforms = None
train_loader = build_dataset(args,
train_df,
transforms=train_transforms,
is_train=True)
batch_num = len(train_loader)
valid_loader = build_dataset(args,
valid_df_sub,
transforms=valid_transforms,
is_train=False)
valid_loader1 = build_dataset(args,
valid_df_sub1,
transforms=valid_transforms,
is_train=False)
valid_loader2 = build_dataset(args,
valid_df_sub2,
transforms=valid_transforms,
is_train=False)
model = build_model(args, device)
if args.model == 'mobilenet_v2':
save_path = os.path.join(args.PATH, 'weights', f'mobilenet_v2_best.pt')
elif args.model == 'resnet18':
save_path = os.path.join(args.PATH, 'weights', f'resnet18_best.pt')
elif args.model == 'resnet50':
save_path = os.path.join(args.PATH, 'weights', f'resnet50_best.pt')
elif args.model == 'resnext':
save_path = os.path.join(args.PATH, 'weights', f'resnext_best.pt')
elif args.model == 'xception':
save_path = os.path.join(args.PATH, 'weights', f'xception_best.pt')
else:
NotImplementedError
if args.model_type == 'lrcn':
save_path = os.path.join(args.PATH, 'weights', f'lrcn_best.pt')
optimizer = build_optimizer(args, model)
scheduler = build_scheduler(args, optimizer, batch_num)
train_cfg = {
'train_loader': train_loader,
'valid_loader': valid_loader,
'valid_loader1': valid_loader1,
'valid_loader2': valid_loader2,
'model': model,
'criterion': nn.BCEWithLogitsLoss(),
'optimizer': optimizer,
'scheduler': scheduler,
'save_path': save_path,
'device': device
}
train_model(args, train_cfg)
def main():
global args, config, best_loss
args = parser.parse_args()
with open(args.config) as f:
config = yaml.load(f)
for k, v in config['common'].items():
setattr(args, k, v)
config = EasyDict(config['common'])
rank, world_size, device_id = dist_init(
os.path.join(args.distributed_path, config.distributed_file))
args.save_path_dated = args.save_path + '/' + args.datetime
if args.run_tag != '':
args.save_path_dated += '-' + args.run_tag
# create model
model = model_entry(config.model)
model.cuda()
model = nn.parallel.DistributedDataParallel(model, device_ids=[device_id])
# create optimizer
opt_config = config.optimizer
opt_config.kwargs.lr = config.lr_scheduler.base_lr
opt_config.kwargs.params = model.parameters()
optimizer = optim_entry(opt_config)
# optionally resume from a checkpoint
last_iter = -1
best_loss = 1e9
if args.load_path:
if args.recover:
best_loss, last_iter = load_state(args.load_path,
model,
optimizer=optimizer)
else:
load_state(args.load_path, model)
cudnn.benchmark = True
# train augmentation
if config.augmentation.get('imgnet_mean', False):
model_mean = (0.485, 0.456, 0.406)
model_std = (0.229, 0.224, 0.225)
else:
model_mean = (0.5, 0.5, 0.5)
model_std = (0.5, 0.5, 0.5)
trans = albumentations.Compose([
RandomResizedCrop(config.augmentation.input_size,
config.augmentation.input_size,
scale=(config.augmentation.min_scale**2., 1.),
ratio=(1., 1.)),
HorizontalFlip(p=0.5),
RandomBrightnessContrast(brightness_limit=0.25,
contrast_limit=0.1,
p=0.5),
JpegCompression(p=.2, quality_lower=50),
MotionBlur(p=0.5),
Normalize(mean=model_mean, std=model_std),
ToTensorV2()
])
train_dataset = FaceDataset(config.train_root,
config.train_source,
transform=trans,
resize=config.augmentation.input_size,
image_format=config.get('image_format', None),
random_frame=config.get(
'train_random_frame', False),
bgr=config.augmentation.get('bgr', False))
train_sampler = DistributedGivenIterationSampler(
train_dataset,
config.lr_scheduler.max_iter,
config.batch_size,
last_iter=last_iter)
train_loader = DataLoader(train_dataset,
batch_size=config.batch_size,
shuffle=False,
num_workers=config.workers,
pin_memory=True,
sampler=train_sampler)
# validation augmentation
trans = albumentations.Compose([
Resize(config.augmentation.input_size, config.augmentation.input_size),
Normalize(mean=model_mean, std=model_std),
ToTensorV2()
])
val_multi_loader = []
if args.val_source != '':
for dataset_idx in range(len(args.val_source)):
val_dataset = FaceDataset(
args.val_root[dataset_idx],
args.val_source[dataset_idx],
transform=trans,
output_index=True,
resize=config.augmentation.input_size,
image_format=config.get('image_format', None),
bgr=config.augmentation.get('bgr', False))
val_sampler = DistributedSampler(val_dataset, round_up=False)
val_loader = DataLoader(val_dataset,
batch_size=config.batch_size,
shuffle=False,
num_workers=config.workers,
pin_memory=True,
sampler=val_sampler)
val_multi_loader.append(val_loader)
config.lr_scheduler['optimizer'] = optimizer
config.lr_scheduler['last_iter'] = last_iter
lr_scheduler = get_scheduler(config.lr_scheduler)
if rank == 0:
mkdir(args.save_path)
mkdir(args.save_path_dated)
tb_logger = SummaryWriter(args.save_path_dated)
logger = create_logger('global_logger',
args.save_path_dated + '-log.txt')
logger.info('{}'.format(args))
logger.info(model)
logger.info(parameters_string(model))
logger.info('len(train dataset) = %d' % len(train_loader.dataset))
for dataset_idx in range(len(val_multi_loader)):
logger.info(
'len(val%d dataset) = %d' %
(dataset_idx, len(val_multi_loader[dataset_idx].dataset)))
mkdir(args.save_path_dated + '/saves')
else:
tb_logger = None
positive_weight = config.get('positive_weight', 0.5)
weight = torch.tensor([1. - positive_weight, positive_weight]) * 2.
if rank == 0:
logger.info('using class weights: {}'.format(weight.tolist()))
criterion = nn.CrossEntropyLoss(weight=weight).cuda()
if args.evaluate:
if args.evaluate_path:
all_ckpt = get_all_checkpoint(args.evaluate_path, args.range_list,
rank)
for ckpt in all_ckpt:
if rank == 0:
logger.info('Testing ckpt: ' + ckpt)
last_iter = -1
_, last_iter = load_state(ckpt, model, optimizer=optimizer)
for dataset_idx in range(len(val_multi_loader)):
validate(dataset_idx,
val_multi_loader[dataset_idx],
model,
criterion,
tb_logger,
curr_step=last_iter,
save_softmax=True)
else:
for dataset_idx in range(len(val_multi_loader)):
validate(dataset_idx,
val_multi_loader[dataset_idx],
model,
criterion,
tb_logger,
curr_step=last_iter,
save_softmax=True)
return
train(train_loader, val_multi_loader, model, criterion, optimizer,
lr_scheduler, last_iter + 1, tb_logger)
return
fake_img = fake_img['image']
real_img = np.rollaxis(real_img, 2, 0)
fake_img = np.rollaxis(fake_img, 2, 0)
return (real_img, torch.tensor(0, dtype=torch.float32)), (fake_img, torch.tensor(1, dtype=torch.float32))
train_transforms = albumentations.Compose([
# ShiftScaleRotate(p=0.3, scale_limit=0.25, border_mode=1, rotate_limit=25),
HorizontalFlip(p=0.3),
# RandomBrightnessContrast(p=0.2, brightness_limit=0.25, contrast_limit=0.5),
# MotionBlur(p=0.2),
GaussNoise(p=0.3),
JpegCompression(p=0.3, quality_lower=50),
Normalize()
])
valid_transforms = albumentations.Compose([
HorizontalFlip(p=0.2),
albumentations.OneOf([
JpegCompression(quality_lower=8, quality_upper=30, p=1.0),
# Downscale(scale_min=0.25, scale_max=0.75, p=1.0),
GaussNoise(p=1.0),
], p=0.22),
Normalize()
])
def build_dataset(args, df, is_train=False):
def get_transforms(config, mode: str = "train") -> Compose:
"""
Composes albumentations transforms.
Returns the full list of transforms when mode is "train".
mode should be one of "train", "val".
"""
# compose validation transforms
if mode == "val":
transforms = Compose(
[],
bbox_params=BboxParams(
format="pascal_voc",
min_area=0.0,
min_visibility=0.0,
label_fields=["category_id"],
),
)
# compose train transforms
# TODO: make transformation parameters configurable from yml
elif mode == "train":
transforms = Compose(
[
LongestMaxSize(
max_size=config["LONGESTMAXSIZE_MAXSIZE"],
p=config["LONGESTMAXSIZE_P"],
),
# PadIfNeeded(min_height=768, min_width=768, border_mode=0, p=1),
RandomSizedBBoxSafeCrop(
height=config["RANDOMSIZEDBBOXSAFECROP_HEIGHT"],
width=config["RANDOMSIZEDBBOXSAFECROP_WIDTH"],
p=config["LONGESTMAXSIZE_P"],
),
ShiftScaleRotate(
shift_limit=config["SHIFTSCALEROTATE_SHIFTLIMIT"],
scale_limit=config["SHIFTSCALEROTATE_SCALELIMIT"],
rotate_limit=config["SHIFTSCALEROTATE_ROTATELIMIT"],
p=config["SHIFTSCALEROTATE_P"],
),
HorizontalFlip(p=config["HORIZONTALFLIP_P"]),
RandomRotate90(p=config["RANDOMROTATE90_P"]),
RandomBrightnessContrast(
brightness_limit=config[
"RANDOMBRIGHTNESSCONTRAST_BRIGHTNESSLIMIT"],
contrast_limit=config[
"RANDOMBRIGHTNESSCONTRAST_CONTRASTLIMIT"],
p=config["RANDOMBRIGHTNESSCONTRAST_P"],
),
RandomGamma(
gamma_limit=config["RANDOMGAMMA_GAMMALIMIT"],
p=config["RANDOMGAMMA_P"],
),
HueSaturationValue(
hue_shift_limit=config["HUESATURATIONVALUE_HUESHIFTLIMIT"],
sat_shift_limit=config["HUESATURATIONVALUE_SATSHIFTLIMIT"],
val_shift_limit=config["HUESATURATIONVALUE_VALSHIFTLIMIT"],
p=config["HUESATURATIONVALUE_P"],
),
MotionBlur(
blur_limit=tuple(config["MOTIONBLUR_BLURLIMIT"]),
p=config["MOTIONBLUR_P"],
),
JpegCompression(
quality_lower=config["JPEGCOMPRESSION_QUALITYLOWER"],
quality_upper=config["JPEGCOMPRESSION_QUALITYUPPER"],
p=config["JPEGCOMPRESSION_P"],
),
],
bbox_params=BboxParams(
format="pascal_voc",
min_area=0.0,
min_visibility=0.0,
label_fields=["category_id"],
),
)
return transforms