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)