from data import get_loader, test_dataset
from options import opt
import torch.backends.cudnn as cudnn
import logging
# set the device for training
if opt.gpu_id == '0':
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
print('USE GPU 0')
elif opt.gpu_id == '1':
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
print('USE GPU 1')
cudnn.benchmark = True
# build the model
model = BBSNet()
if (opt.load is not None):
model.load_state_dict(torch.load(opt.load))
print('load model from ', opt.load)
model.cuda()
params = model.parameters()
optimizer = torch.optim.Adam(params, opt.lr)
# set the path
image_root = opt.rgb_root
gt_root = opt.gt_root
depth_root = opt.depth_root
test_image_root = opt.test_rgb_root
test_gt_root = opt.test_gt_root
test_depth_root = opt.test_depth_root
default='../BBS_dataset/RGBD_for_test/',
help='test dataset path')
opt = parser.parse_args()
dataset_path = opt.test_path
#set device for test
if opt.gpu_id == '0':
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
print('USE GPU 0')
elif opt.gpu_id == '1':
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
print('USE GPU 1')
#load the model
model = BBSNet()
#Large epoch size may not generalize well. You can choose a good model to load according to the log file and pth files saved in ('./BBSNet_cpts/') when training.
model.load_state_dict(torch.load('./model_pths/BBSNet.pth'))
model.cuda()
model.eval()
#test
test_datasets = ['NJU2K', 'NLPR', 'STERE', 'DES', 'SSD', 'LFSD', 'SIP']
for dataset in test_datasets:
save_path = './test_maps/BBSNet/ResNet50/' + dataset + '/'
if not os.path.exists(save_path):
os.makedirs(save_path)
image_root = dataset_path + dataset + '/RGB/'
gt_root = dataset_path + dataset + '/GT/'
depth_root = dataset_path + dataset + '/depth/'
test_loader = test_dataset(image_root, gt_root, depth_root, opt.testsize)
default='./datasets/',
help='test dataset path')
opt = parser.parse_args()
dataset_path = opt.test_path
# set device for test
if opt.gpu_id == '0':
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
print('USE GPU 0')
elif opt.gpu_id == '1':
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
print('USE GPU 1')
# load the model
model = BBSNet()
# Large epoch size may not generalize well. You can choose a good model to load according to the log file and pth files saved in ('./BBSNet_cpts/') when training.
model.load_state_dict(torch.load('./eval/BBS-Net/BBSNet_epoch_best.pth'))
model.cuda()
model.eval()
# test
test_datasets = [
'STERE', 'NJU2K_TEST', 'NLPR_TEST', 'DES', 'SSD', 'LFSD', 'SIP'
]
config = init_wnb_config()
with wandb.init(project=config['wandb_project'], config=config):
for dataset in test_datasets:
print(f'Processing {dataset}')
if not os.path.exists(dataset_path + dataset):
print(f'Dataset {dataset_path + dataset} does not exist')
default='../BBS_dataset/RGBD_for_test/',
help='test dataset path')
opt = parser.parse_args()
dataset_path = opt.test_path
#set device for test
if opt.gpu_id == '0':
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
print('USE GPU 0')
elif opt.gpu_id == '1':
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
print('USE GPU 1')
#load the model
model = BBSNet()
#Large epoch size may not generalize well. You can choose a good model to load according to the log file and pth files saved in ('./BBSNet_cpts/') when training.
model.load_state_dict(torch.load('./BBSNet_cpts/BBSNet_epoch_195.pth')
) #'./model_pths/BBSNet_epoch_best.pth'))
model.cuda()
model.eval()
#test
test_datasets = ['NJU2K', 'NLPR', 'STERE', 'DES', 'SSD', 'LFSD', 'SIP']
for dataset in test_datasets:
save_path = './test_maps/BBSNet/epoch195/' + dataset + '/'
if not os.path.exists(save_path):
os.makedirs(save_path)
image_root = dataset_path + dataset + '/RGB/'
gt_root = dataset_path + dataset + '/GT/'
depth_root = dataset_path + dataset + '/depth/'