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)
