-
GitHub链接
https://github.com/pkuCactus/BDCN
-
环境配置(Prerequisites)
Python 3.8
contourpy==1.1.1
cycler==0.12.1
filelock==3.14.0
fonttools==4.53.0
fsspec==2024.6.0
importlib_resources==6.4.0
intel-openmp==2021.4.0
Jinja2==3.1.4
kiwisolver==1.4.5
MarkupSafe==2.1.5
matplotlib==3.7.5
mkl==2021.4.0
mpmath==1.3.0
networkx==3.1
numpy==1.24.4
opencv-python==4.10.0.82
packaging==24.1
pillow==10.3.0
pyparsing==3.1.2
python-dateutil==2.9.0.post0
scipy==1.10.1
six==1.16.0
sympy==1.12.1
tbb==2021.12.0
torch==2.3.1
torchvision==0.18.1
typing_extensions==4.12.2
zipp==3.19.2
-
准备工作
-
下载预训练模型
https://pan.baidu.com/s/10Tgjs7FiAYWjVyVgvEM0mA
code: ab4g
-
新建一个文件夹存放测试图片
-
-
新建结果存放文件夹
-
代码修改
test_image.py：
cpu:添加map_location=torch.device('cpu')
gpu:去掉map_location=torch.device('cpu')
cv2.imwrite(os.path.join(save_dir, 's2d_' + str(k), '%s' % nm), 255 - out[j] * 255)
cv2.imwrite(os.path.join(save_dir, 'd2s_' + str(k), '%s' % nm), 255 - 255 * out[j+5])
p.s. 图片分辨率过高，会非常消耗内存
-
device问题：model.load_state_dict(torch.load('%s' % (args.model), map_location=torch.device('cpu')))。
-
args.rate: 去掉model = bdcn.BDCN(args.rate)
-
save_dir = args.res_dir
if not os.path.exists(save_dir):
os.mkdir(save_dir)移到前面
-
43行左右，添加k=1，且在循环中添加k+=1
-
data = cv2.imread(test_root + '/' + nm+'.jpg')去掉'.jpg'
-
80行左右修改：
-
-
运行代码
def parse_args(): parser = argparse.ArgumentParser('test BDCN') parser.add_argument('-c', '--cuda', action='store_true', help='whether use gpu to train network') parser.add_argument('-g', '--gpu', type=str, default='0', help='the gpu id to train net') parser.add_argument('-m', '--model', type=str, default='bdcn-final-model/bdcn_pretrained_on_bsds500.pth', help='the model to test') parser.add_argument('--res-dir', type=str, default='result', # folder of results help='the dir to store result') parser.add_argument('--data-root', type=str, default='testdata') # 测试文件夹 folder of test images parser.add_argument('--test-lst', type=str, default=None) return parser.parse_args()
-
数据集
BSDS:http://mftp.mmcheng.net/liuyun/rcf/data/HED-BSDS.tar.gz
-
结果图
d2s_1
s2d_1
d2s_2
s2d_2
d2s_3
s2d_3
d2s_4
s2d_4
d2s_5
s2d_5
fuse
-
test_image.py文件
import numpy as np import torch import torch.optim as optim import torch.nn as nn from torch.autograd import Variable from torch.nn import functional as F import time import re import os import sys import cv2 import bdcn from datasets.dataset import Data import argparse import cfg from matplotlib import pyplot as plt import os import os.path as osp from scipy.io import savemat def make_dir(data_dir): if not os.path.exists(data_dir): os.mkdir(data_dir) def test(model, args): test_root = args.data_root if args.test_lst is not None: with open(osp.join(test_root, args.test_lst), 'r') as f: test_lst = f.readlines() test_lst = [x.strip() for x in test_lst] if ' ' in test_lst[0]: test_lst = [x.split(' ')[0] for x in test_lst] else: test_lst = os.listdir(test_root) print(test_lst[0]) save_dir = args.res_dir if not os.path.exists(save_dir): os.mkdir(save_dir) k = 1 save_sideouts = 1 if save_sideouts: for j in range(5): make_dir(os.path.join(save_dir, 's2d_' + str(k))) make_dir(os.path.join(save_dir, 'd2s_' + str(k))) k += 1 mean_bgr = np.array([104.00699, 116.66877, 122.67892]) if args.cuda: model.cuda() model.eval() start_time = time.time() all_t = 0 for nm in test_lst: data = cv2.imread(test_root + '/' + nm) print(f'data:{data}') # print(os.path.join(test_root, nm)) # data = cv2.resize(data, (data.shape[1]/2, data.shape[0]/2), interpolation=cv2.INTER_LINEAR) data = np.array(data, np.float32) data -= mean_bgr data = data.transpose((2, 0, 1)) data = torch.from_numpy(data).float().unsqueeze(0) if args.cuda: data = data.cuda() data = Variable(data) t1 = time.time() out = model(data) if '/' in nm: nm = nm.split('/')[-1] print("nm:", nm) if save_sideouts: out = [F.sigmoid(x).cpu().data.numpy()[0, 0, :, :] for x in out] print(f'out:{len(out)}') k = 1 for j in range(5): # savemat(osp.join(save_dir, 's2d_'+str(k), nm+'.mat'), {'prob': out[j]}) cv2.imwrite(os.path.join(save_dir, 's2d_' + str(k), '%s' % nm), 255 - out[j] * 255) # savemat(osp.join(save_dir, 'd2s_'+str(k), nm+'.mat'), {'prob': out[j+5]}) cv2.imwrite(os.path.join(save_dir, 'd2s_' + str(k), '%s' % nm), 255 - 255 * out[j+5]) k += 1 else: out = [F.sigmoid(out[-1]).cpu().data.numpy()[0, 0, :, :]] if not os.path.exists(os.path.join(save_dir, 'fuse')): os.mkdir(os.path.join(save_dir, 'fuse')) cv2.imwrite(os.path.join(save_dir, 'fuse/%s.png' % nm.split('/')[-1].split('.')[0]), 255 * out[-1]) all_t += time.time() - t1 print(all_t) print('Overall Time use: ', time.time() - start_time) def main(): import time print(time.localtime()) args = parse_args() os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu # model = bdcn.BDCN(rate=args.rate) model = bdcn.BDCN() model.load_state_dict(torch.load('%s' % (args.model), map_location=torch.device('cpu'))) # print model.fuse.weight.data, model.fuse.bias.data print(model.fuse.weight.data) test(model, args) def parse_args(): parser = argparse.ArgumentParser('test BDCN') parser.add_argument('-c', '--cuda', action='store_true', help='whether use gpu to train network') parser.add_argument('-g', '--gpu', type=str, default='0', help='the gpu id to train net') parser.add_argument('-m', '--model', type=str, default='bdcn-final-model/bdcn_pretrained_on_bsds500.pth', help='the model to test') parser.add_argument('--res-dir', type=str, default='result', help='the dir to store result') parser.add_argument('--data-root', type=str, default='testdata') parser.add_argument('--test-lst', type=str, default=None) return parser.parse_args() if __name__ == '__main__': main()
