def main():
global args, best_prec1, model, train_dataset, val_dataset
args = parser.parse_args()
logging.basicConfig(
format="%(message)s",
handlers=[
logging.FileHandler("{0}/{1}.log".format(args.log, sys.argv[0].replace('.py','') + datetime.now().strftime('_%H_%M_%d_%m_%Y'))),
logging.StreamHandler()
],
level=logging.INFO)
# create model
# if args.pretrained:
# pass
# # logging.info("=> using pre-trained model '{}'".format(args.arch))
# # model = models.__dict__[args.arch](pretrained=True)
# else:
# logging.info("=> creating model '{}'".format(args.arch))
model = SpatioTemporalSaliency(CONFIG)
model._initialize_weights()
model.eval()
# define loss function (criterion) and optimizer
criterion = nn.BCELoss().cuda()
for param in model.parameters():
param.requires_grad = False
# optionally resume from a checkpoint
cudnn.benchmark = True
# Data loading code
config = CONFIG.copy()
epochs = args.epochs
users = 15
len_dataset = 700
masks = np.zeros((epochs, users, len_dataset, 8, 75, 100), dtype=np.float16)
bad_list = set([231,310,372,317,436,441,447,541,535,675,661])
for epoch in range(1,epochs+1):
print('starting forward pass for epoch {0}'.format(epoch))
for user in CONFIG['test']:
print('starting user {0}'.format(user+1))
config['test'] = [user]
test_dataset = SequnceDataset(config, 'train')
# Let's resume weights.
w_path = os.path.join(args.weights, '{0}_{1}.pth.tar'.format(user, epoch))
if os.path.isfile(w_path):
logging.info("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(w_path)
args.start_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['state_dict'])
logging.info("=> loaded checkpoint '{}' (epoch {})"
.format(args.resume, checkpoint['epoch']))
else:
logging.info("=> no checkpoint found at '{}'".format(args.resume))
start = time.time()
for img_idx, (input, sal, target, img_path) in enumerate(test_dataset):
img_idx = int(img_path.split('/')[-1].split('.jpg')[0]) - 1001
print(epoch, user, img_idx, len(test_dataset))
if img_idx in bad_list:
continue
# measure data loading time
input_var = torch.autograd.Variable(input, volatile=True).cuda(0)
output = model([input_var, sal, target, img_path])
masks[epoch-1][user][img_idx] = output[0,0]
img = Image.open(img_path)
w, h = img.size
len_out = output.shape[0]
counter = 0
path = os.path.join(args.visualize, str(user), str(epoch))
if not os.path.exists(path):
os.makedirs(path)
for seq_idx, tar in enumerate(output):
if seq_idx >= len_out:
break
mask = np.array(output[seq_idx][0] * 255, dtype=np.uint8)
mask = Image.fromarray(mask).resize((w,h)).convert('RGB')
saliency = np.array(tar * 255, dtype=np.uint8)
saliency = Image.fromarray(saliency).resize((w,h)).convert('RGB')
out = Image.new('RGB', (w, h*2))
out.paste(Image.blend(img, mask, alpha=0.7).convert('RGB'), (0,0))
out.paste(Image.blend(img, saliency, alpha=0.7).convert('RGB'),(0,h))
out_path = os.path.join(path, '{0}-{1}.jpg'.format(img_idx, seq_idx))
out.save(out_path)
print(time.time()-start)
def main():
global args, best_prec1, model, train_dataset, val_dataset
args = parser.parse_args()
#
# # Data loading code
config = CONFIG.copy()
mode = args.mode
for model_name in MODELS:
# setting up the logger.
log = logging.getLogger()
for hdlr in log.handlers[:]:
if isinstance(hdlr, logging.FileHandler):
log.removeHandler(hdlr)
logging.basicConfig(
format="%(message)s",
handlers=[
logging.FileHandler("{0}/{1}-{2}.log".format(
args.log, model_name, sys.argv[0].replace('.py', '') +
datetime.now().strftime('_%H_%M_%d_%m_%Y'))),
logging.StreamHandler()
],
level=logging.INFO)
for user in CONFIG[mode]['users']:
config[mode]['users'] = [user]
config['model']['name'] = model_name
train_dataset = SequnceDataset(config, mode)
#logging.info("=> creating model'{}'".format(args.arch))
if config['model']['type'] == 'RNN':
model = RNNSaliency(config)
else:
pass
model._initialize_weights(pretrained=True)
# model = model.cuda()
for param in model.encoder.parameters():
param.requires_grad = False
criterion = nn.BCELoss().cuda()
# criterion = nn.MSELoss().cuda()
optimizer = torch.optim.Adam(
# list(model.readout.parameters()) +\
model.decoder.parameters(),
config['train']['lr'],
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=config['train']['weight_decay'])
# optimizer = torch.optim.SGD(
# # list(model.readout.parameters()) +\
# list(model.decoder.parameters()),
# config['train']['lr'],
# momentum=args.momentum,
# weight_decay=config['train']['weight_decay'])
# if args.resume:
# if os.path.isfile(args.resume):
# logging.info("=> loading checkpoint '{}'".format(args.resume))
# checkpoint = torch.load(args.resume)
# args.start_epoch = checkpoint['epoch']
# # best_prec1 = checkpoint['best_prec1']
# model.load_state_dict(checkpoint['state_dict'])
# optimizer.load_state_dict(checkpoint['optimizer'])
# logging.info("=> loaded checkpoint '{}' (epoch {})"
# .format(args.resume, checkpoint['epoch']))
# else:
# logging.info("=> no checkpoint found at '{}'".format(args.resume))
# train_loader = train_dataset
# train_loader = torch.utils.data.DataLoader(
# train_dataset, batch_size=args.batch_size, shuffle=False,
# num_workers=args.workers, pin_memory=True, sampler=None)
# val_loader = torch.utils.data.DataLoader(
# Saliency( CONFIG, 'test'),
# batch_size=args.batch_size, shuffle=False,
# num_workers=args.workers, pin_memory=True)
# if args.evaluate:
# validate(val_loader, model, criterion)
# return
# if args.visualize:
# visualize(train_loader, model, user, epoch)
# return
for epoch in range(args.start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch)
# train for one epoch
train(train_dataset, model, criterion, optimizer, epoch,
config)
# evaluate on validation set
#if epoch%10 == 0:
validate(train_dataset, model, criterion, user, epoch, config)
# visualize(train_loader, model, str(user), str(epoch))
# remember best prec@1 and save checkpoint
# is_best = prec1 > best_prec1
# best_prec1 = max(prec1, best_prec1)
save_checkpoint({
'epoch': epoch + 1,
'user': user,
'arch': config['model']['name'],
'state_dict': model.state_dict(),
# 'best_prec1': best_prec1,
# 'optimizer' : optimizer.state_dict(),
})
