Spaces:
Sleeping
Sleeping
File size: 4,306 Bytes
68c537d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 |
import argparse
from utils import *
import torchvision.transforms as transforms
from net.dornet import Net
from torch.utils.data import Dataset, DataLoader
from data.nyu_dataloader import *
from data.rgbdd_dataloader import *
# from data.tofsr_dataloader import *
import os
import torch
parser = argparse.ArgumentParser()
parser.add_argument('--scale', type=int, default=4, help='scale factor')
parser.add_argument("--root_dir", type=str, default='./dataset/RGB-D-D', help="root dir of dataset")
parser.add_argument("--model_dir", type=str, default="./checkpoints/RGBDD.pth", help="path of net")
parser.add_argument("--results_dir", type=str, default='./results/', help="root dir of results")
parser.add_argument('--tiny_model', action='store_true', help='tiny model')
parser.add_argument("--blur_sigma", type=int, default=3.6, help="blur_sigma")
parser.add_argument('--isNoisy', action='store_true', help='Noisy')
opt = parser.parse_args()
net = Net(tiny_model=True).cuda()
print("*********************************************")
print(sum(p.numel() for p in net.parameters() if p.requires_grad))
print("*********************************************")
net.load_state_dict(torch.load(opt.model_dir, map_location='cuda:0'))
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net.to(device)
data_transform = transforms.Compose([transforms.ToTensor()])
dataset_name = opt.root_dir.split('/')[-1]
if dataset_name == 'RGB-D-D':
dataset = RGBDD_Dataset(root_dir=opt.root_dir, scale=opt.scale, downsample='real', train=False,
transform=data_transform, isNoisy=opt.isNoisy, blur_sigma=opt.blur_sigma)
rmse = np.zeros(405)
elif dataset_name == 'NYU-v2':
dataset = NYU_v2_datset(root_dir=opt.root_dir, scale=opt.scale, transform=data_transform, train=False)
test_minmax = np.load('%s/test_minmax.npy' % opt.root_dir)
rmse = np.zeros(449)
dataloader = DataLoader(dataset, batch_size=1, shuffle=False, num_workers=8)
data_num = len(dataloader)
with torch.no_grad():
net.eval()
if dataset_name == 'RGB-D-D':
for idx, data in enumerate(dataloader):
guidance, lr, gt, maxx, minn, name = data['guidance'].cuda(), data['lr'].cuda(), data['gt'].cuda(), data[
'max'].cuda(), data['min'].cuda(), data['name']
out = net(x_query=lr, rgb=guidance)
rmse[idx] = rgbdd_calc_rmse(gt[0, 0], out[0, 0], [maxx, minn])
# Save results (Save the output depth map)
# path_output = '{}/output'.format(opt.results_dir)
# os.makedirs(path_output, exist_ok=True)
# path_save_pred = '{}/{}.png'.format(path_output, name[0])
# pred = out[0, 0] * (maxx - minn) + minn
# pred = pred.cpu().detach().numpy()
# pred = pred.astype(np.uint16)
# pred = Image.fromarray(pred)
# pred.save(path_save_pred)
print('idx:%d RMSE:%f' % (idx + 1, rmse[idx]))
print("==========RGB-D-D=========")
print(rmse.mean())
print("==========RGB-D-D=========")
elif dataset_name == 'NYU-v2':
# t = np.zeros(449)
for idx, data in enumerate(dataloader):
guidance, lr, gt = data['guidance'].cuda(), data['lr'].cuda(), data['gt'].cuda()
out = net(x_query=lr, rgb=guidance)
minmax = test_minmax[:, idx]
minmax = torch.from_numpy(minmax).cuda()
rmse[idx] = calc_rmse(gt[0, 0], out[0, 0], minmax)
# Save results (Save the output depth map)
# path_output = '{}/output'.format(opt.results_dir)
# os.makedirs(path_output, exist_ok=True)
# path_save_pred = '{}/{:010d}.png'.format(path_output, idx)
# pred = out[0,0] * (minmax[0] - minmax[1]) + minmax[1]
# pred = pred * 1000.0
# pred = pred.cpu().detach().numpy()
# pred = pred.astype(np.uint16)
# pred = Image.fromarray(pred)
# pred.save(path_save_pred)
print('idx:%d RMSE:%f' % (idx + 1, rmse[idx]))
print("=========NYU-v2==========")
print(rmse.mean())
print("=========NYU-v2==========")
|