bdcn手顺

1. GitHub链接

https://github.com/pkuCactus/BDCN

1. 环境配置(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

2. 准备工作

   

   1. 下载预训练模型

          https://pan.baidu.com/s/10Tgjs7FiAYWjVyVgvEM0mA

          code: ab4g

   2. 新建一个文件夹存放测试图片

• 新建结果存放文件夹

1. 代码修改

     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. 图片分辨率过高，会非常消耗内存

   1. device问题：model.load_state_dict(torch.load('%s' % (args.model), map_location=torch.device('cpu')))。

   2. args.rate: 去掉model = bdcn.BDCN(args.rate)

   3. save_dir = args.res_dir

          if not os.path.exists(save_dir):

          os.mkdir(save_dir)移到前面

   4. 43行左右，添加k=1，且在循环中添加k+=1

   5. data = cv2.imread(test_root + '/' + nm+'.jpg')去掉'.jpg'

   6. 80行左右修改：

2. 运行代码

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()

1. 数据集

BSDS:http://mftp.mmcheng.net/liuyun/rcf/data/HED-BSDS.tar.gz

1. 结果图

d2s_1

s2d_1

d2s_2

s2d_2

d2s_3

s2d_3

d2s_4

s2d_4

d2s_5

s2d_5

fuse

1. 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()