def test_one_image(args, dt_config, dataset_class):
input_size = (475, 475)
model_path = args.snapshot
dataset_instance = dataset_class(data_path=dt_config.DATA_PATH)
num_classes = dataset_instance.num_classes
model = PSPNet(num_classes=num_classes)
model.load_state_dict(torch.load(model_path)["state_dict"])
model.eval()
img = cv2.imread(args.image_path)
processed_img = cv2.resize(img, input_size)
overlay = np.copy(processed_img)
processed_img = processed_img / 255.0
processed_img = torch.tensor(
processed_img.transpose(2, 0, 1)[np.newaxis, :]).float()
if torch.cuda.is_available():
model = model.cuda()
processed_img = processed_img.cuda()
output = model(processed_img)[0]
mask = output.data.max(1)[1].cpu().numpy().reshape(475, 475)
color_mask = np.array(dataset_instance.colors)[mask]
alpha = args.alpha
overlay = (((1 - alpha) * overlay) + (alpha * color_mask)).astype("uint8")
overlay = cv2.resize(overlay, (img.shape[1], img.shape[0]))
cv2.imwrite("result.jpg", overlay)
def main():
args = parse_arguments()
# Dataset used for training the model
MEAN = [0.45734706, 0.43338275, 0.40058118]
STD = [0.23965294, 0.23532275, 0.2398498]
to_tensor = transforms.ToTensor()
normalize = transforms.Normalize(MEAN, STD)
num_classes = 2
palette = [0, 0, 0, 128, 0, 128]
# Model
model = PSPNet(num_classes=num_classes, backbone='resnet18')
availble_gpus = list(range(torch.cuda.device_count()))
device = torch.device('cuda:0' if len(availble_gpus) > 0 else 'cpu')
checkpoint = torch.load(args.model)
if isinstance(checkpoint, dict) and 'state_dict' in checkpoint.keys():
checkpoint = checkpoint['state_dict']
if 'module' in list(checkpoint.keys())[0] and not isinstance(
model, torch.nn.DataParallel):
model = torch.nn.DataParallel(model)
model.load_state_dict(checkpoint)
model.to(device)
model.eval()
if not os.path.exists('outputs'):
os.makedirs('outputs')
image_files = sorted(glob(os.path.join(args.images,
f'*.{args.extension}')))
with torch.no_grad():
tbar = tqdm(image_files, ncols=100)
for img_file in tbar:
image = Image.open(img_file).convert('RGB')
image = image.resize((480, 320))
input = normalize(to_tensor(image)).unsqueeze(0)
print(input.size())
t1 = time.time()
prediction = model(input.to(device))
prediction = prediction.squeeze(0).cpu().numpy()
print(time.time() - t1)
prediction = F.softmax(torch.from_numpy(prediction),
dim=0).argmax(0).cpu().numpy()
save_images(image, prediction, args.output, img_file, palette)
def main():
batch_size = 8
net = PSPNet(pretrained=False, num_classes=num_classes, input_size=(512, 1024)).cuda()
snapshot = 'epoch_48_validation_loss_5.1326_mean_iu_0.3172_lr_0.00001000.pth'
net.load_state_dict(torch.load(os.path.join(ckpt_path, snapshot)))
net.eval()
mean_std = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
transform = transforms.Compose([
expanded_transform.FreeScale((512, 1024)),
transforms.ToTensor(),
transforms.Normalize(*mean_std)
])
restore = transforms.Compose([
expanded_transform.DeNormalize(*mean_std),
transforms.ToPILImage()
])
lsun_path = '/home/b3-542/LSUN'
dataset = LSUN(lsun_path, ['tower_val', 'church_outdoor_val', 'bridge_val'], transform=transform)
dataloader = DataLoader(dataset, batch_size=batch_size, num_workers=16, shuffle=True)
if not os.path.exists(test_results_path):
os.mkdir(test_results_path)
for vi, data in enumerate(dataloader, 0):
inputs, labels = data
inputs = Variable(inputs, volatile=True).cuda()
outputs = net(inputs)
prediction = outputs.cpu().data.max(1)[1].squeeze_(1).numpy()
for idx, tensor in enumerate(zip(inputs.cpu().data, prediction)):
pil_input = restore(tensor[0])
pil_output = colorize_mask(tensor[1])
pil_input.save(os.path.join(test_results_path, '%d_img.png' % (vi * batch_size + idx)))
pil_output.save(os.path.join(test_results_path, '%d_out.png' % (vi * batch_size + idx)))
print 'save the #%d batch, %d images' % (vi + 1, idx + 1)
def main():
net = PSPNet(num_classes=num_classes)
if len(args['snapshot']) == 0:
# net.load_state_dict(torch.load(os.path.join(ckpt_path, 'cityscapes (coarse)-psp_net', 'xx.pth')))
curr_epoch = 1
args['best_record'] = {'epoch': 0, 'iter': 0, 'val_loss': 1e10, 'acc': 0, 'acc_cls': 0, 'mean_iu': 0,
'fwavacc': 0}
else:
print('training resumes from ' + args['snapshot'])
net.load_state_dict(torch.load(os.path.join(ckpt_path, exp_name, args['snapshot'])))
split_snapshot = args['snapshot'].split('_')
curr_epoch = int(split_snapshot[1]) + 1
args['best_record'] = {'epoch': int(split_snapshot[1]), 'iter': int(split_snapshot[3]),
'val_loss': float(split_snapshot[5]), 'acc': float(split_snapshot[7]),
'acc_cls': float(split_snapshot[9]),'mean_iu': float(split_snapshot[11]),
'fwavacc': float(split_snapshot[13])}
net.cuda().train()
mean_std = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
train_joint_transform = joint_transforms.Compose([
joint_transforms.Scale(args['longer_size']),
joint_transforms.RandomRotate(10),
joint_transforms.RandomHorizontallyFlip()
])
sliding_crop = joint_transforms.SlidingCrop(args['crop_size'], args['stride_rate'], ignore_label)
train_input_transform = standard_transforms.Compose([
standard_transforms.ToTensor(),
standard_transforms.Normalize(*mean_std)
])
val_input_transform = standard_transforms.Compose([
standard_transforms.ToTensor(),
standard_transforms.Normalize(*mean_std)
])
target_transform = extended_transforms.MaskToTensor()
visualize = standard_transforms.Compose([
standard_transforms.Scale(args['val_img_display_size']),
standard_transforms.ToTensor()
])
train_set = Retinaimages('training', joint_transform=train_joint_transform, sliding_crop=sliding_crop,
transform=train_input_transform, target_transform=target_transform)
train_loader = DataLoader(train_set, batch_size=args['train_batch_size'], num_workers=2, shuffle=True)
val_set = Retinaimages('validate', transform=val_input_transform, sliding_crop=sliding_crop,
target_transform=target_transform)
val_loader = DataLoader(val_set, batch_size=1, num_workers=2, shuffle=False)
criterion = CrossEntropyLoss2d(size_average=True).cuda()
optimizer = optim.SGD([
{'params': [param for name, param in net.named_parameters() if name[-4:] == 'bias'],
'lr': 2 * args['lr']},
{'params': [param for name, param in net.named_parameters() if name[-4:] != 'bias'],
'lr': args['lr'], 'weight_decay': args['weight_decay']}
], momentum=args['momentum'], nesterov=True)
if len(args['snapshot']) > 0:
optimizer.load_state_dict(torch.load(os.path.join(ckpt_path, exp_name, 'opt_' + args['snapshot'])))
optimizer.param_groups[0]['lr'] = 2 * args['lr']
optimizer.param_groups[1]['lr'] = args['lr']
check_mkdir(ckpt_path)
check_mkdir(os.path.join(ckpt_path, exp_name))
open(os.path.join(ckpt_path, exp_name, "_1" + '.txt'), 'w').write(str(args) + '\n\n')
train(train_loader, net, criterion, optimizer, curr_epoch, args, val_loader, visualize, val_set)
def main():
net = PSPNet(num_classes=num_classes,
input_size=train_args['input_size']).cuda()
if len(train_args['snapshot']) == 0:
curr_epoch = 0
else:
print 'training resumes from ' + train_args['snapshot']
net.load_state_dict(
torch.load(
os.path.join(ckpt_path, exp_name, train_args['snapshot'])))
split_snapshot = train_args['snapshot'].split('_')
curr_epoch = int(split_snapshot[1])
train_record['best_val_loss'] = float(split_snapshot[3])
train_record['corr_mean_iu'] = float(split_snapshot[6])
train_record['corr_epoch'] = curr_epoch
net.train()
mean_std = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
train_simul_transform = simul_transforms.Compose([
simul_transforms.Scale(int(train_args['input_size'][0] / 0.875)),
simul_transforms.RandomCrop(train_args['input_size']),
simul_transforms.RandomHorizontallyFlip()
])
val_simul_transform = simul_transforms.Compose([
simul_transforms.Scale(int(train_args['input_size'][0] / 0.875)),
simul_transforms.CenterCrop(train_args['input_size'])
])
img_transform = standard_transforms.Compose([
standard_transforms.ToTensor(),
standard_transforms.Normalize(*mean_std)
])
target_transform = standard_transforms.Compose([
expanded_transforms.MaskToTensor(),
expanded_transforms.ChangeLabel(ignored_label, num_classes - 1)
])
restore_transform = standard_transforms.Compose([
expanded_transforms.DeNormalize(*mean_std),
standard_transforms.ToPILImage()
])
train_set = CityScapes('train',
simul_transform=train_simul_transform,
transform=img_transform,
target_transform=target_transform)
train_loader = DataLoader(train_set,
batch_size=train_args['batch_size'],
num_workers=16,
shuffle=True)
val_set = CityScapes('val',
simul_transform=val_simul_transform,
transform=img_transform,
target_transform=target_transform)
val_loader = DataLoader(val_set,
batch_size=val_args['batch_size'],
num_workers=16,
shuffle=False)
weight = torch.ones(num_classes)
weight[num_classes - 1] = 0
criterion = CrossEntropyLoss2d(weight).cuda()
# don't use weight_decay for bias
optimizer = optim.SGD([{
'params': [
param for name, param in net.named_parameters()
if name[-4:] == 'bias' and (
'ppm' in name or 'final' in name or 'aux_logits' in name)
],
'lr':
2 * train_args['new_lr']
}, {
'params': [
param for name, param in net.named_parameters()
if name[-4:] != 'bias' and (
'ppm' in name or 'final' in name or 'aux_logits' in name)
],
'lr':
train_args['new_lr'],
'weight_decay':
train_args['weight_decay']
}, {
'params': [
param
for name, param in net.named_parameters() if name[-4:] == 'bias'
and not ('ppm' in name or 'final' in name or 'aux_logits' in name)
],
'lr':
2 * train_args['pretrained_lr']
}, {
'params': [
param
for name, param in net.named_parameters() if name[-4:] != 'bias'
and not ('ppm' in name or 'final' in name or 'aux_logits' in name)
],
'lr':
train_args['pretrained_lr'],
'weight_decay':
train_args['weight_decay']
}],
momentum=0.9,
nesterov=True)
if len(train_args['snapshot']) > 0:
optimizer.load_state_dict(
torch.load(os.path.join(ckpt_path,
'opt_' + train_args['snapshot'])))
optimizer.param_groups[0]['lr'] = 2 * train_args['new_lr']
optimizer.param_groups[1]['lr'] = train_args['new_lr']
optimizer.param_groups[2]['lr'] = 2 * train_args['pretrained_lr']
optimizer.param_groups[3]['lr'] = train_args['pretrained_lr']
if not os.path.exists(ckpt_path):
os.mkdir(ckpt_path)
if not os.path.exists(os.path.join(ckpt_path, exp_name)):
os.mkdir(os.path.join(ckpt_path, exp_name))
for epoch in range(curr_epoch, train_args['epoch_num']):
train(train_loader, net, criterion, optimizer, epoch)
validate(val_loader, net, criterion, optimizer, epoch,
restore_transform)