shuffle=True,
num_workers=8)
vtransforms = transforms.Compose(
[tr.Resize([img_rows, img_cols]),
tr.ToTensor()])
vdataset = DAVIS2016(root='/home/wilson/DL/project/dataset/davis.pkl',
split='val',
transform=vtransforms)
vdataloader = torch.utils.data.DataLoader(vdataset,
batch_size=1,
shuffle=False,
num_workers=8)
# Setup Metrics
running_metrics = runningScore(pspnet_specs['n_classes'])
# setup Model
base_net = BaseNet()
class_net = ClassNet()
base_net.cuda()
class_net.cuda()
base_opt = torch.optim.Adam(base_net.parameters(), lr=lr, betas=(0.5, 0.999))
class_opt = torch.optim.Adam(class_net.parameters(), lr=lr, betas=(0.5, 0.999))
loss_fn = cross_entropy2d
iter = 0
writer = SummaryWriter('runs/test')
batch_size=1,
shuffle=False,
num_workers=4,
pin_memory=True)
upsample_1024 = nn.Upsample(size=[1024, 2048], mode='bilinear')
model_dict = {}
enc_shared = SharedEncoder().cuda(args.gpu)
model_dict['enc_shared'] = enc_shared
load_models(model_dict, args.weight_dir)
enc_shared.eval()
cty_running_metrics = runningScore(num_classes)
for i_test, (images_test, name) in tqdm(enumerate(test_loader)):
images_test = Variable(images_test.cuda(), volatile=True)
_, _, pred, _ = enc_shared(images_test)
pred = upsample_1024(pred)
pred = pred.data.cpu().numpy()[0]
pred = pred.transpose(1, 2, 0)
pred = np.asarray(np.argmax(pred, axis=2), dtype=np.uint8)
pred = np.asarray(test_set.convert_back_to_id(pred), dtype=np.uint8)
pred = Image.fromarray(pred)
name = name[0][0].split('/')[-1]
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)