Edge detection - PidiNet network trains its own data set and optimizes inference testing (detailed graphic tutorial)

PiDiNet is an algorithm for edge detection that proposes a simple, lightweight yet effective architecture. PiDiNet uses a new

Ying's pixel difference convolution integrates traditional edge detection operators into the popular convolution operations in modern CNN to enhance task performance.

Extensive experiments are conducted on BSDS500, NYUD and Multicue to demonstrate its effectiveness, high training and inference efficiency.

1. Source code package

I have modified some codes myself based on the source code package on the official website. The source code package I provided has added network cropping optimization, CPU inference test code, edge detection map and original image fusion code, etc. It also includes a training set, which is recommended for scholars to download. Use the source code package I provided.

Official website address: PidiNet

The source code package download link I provided: network disk link , extraction code: kmxb

Paper address: paper

The source code package I provided looks like this after decompression:

2. Data set preparation

The enhanced data sets include BSD 500, PASCAL VOC and NYUD. The download link is as follows. Copy the link directly to Thunder and download it quickly. There is a BSD 500 data set in the source code package I provided, located in the directory pidinet\path\to, as follows:

If you have made your own data set, you can directly import it for use.

BSD 500数据集下载链接：   http://mftp.mmcheng.net/liuyun/rcf/data/HED-BSDS.tar.gz
PASCAL VOC数据集下载链接： http://mftp.mmcheng.net/liuyun/rcf/data/PASCAL.tar.gz
NYUD数据集下载链接  http://mftp.mmcheng.net/liuyun/rcf/data/NYUD.tar.gz

3. Training

There are many options to choose from, the official website provides the following:

Among them, the best one is table7_pidinet, and the lightest one is table5_pidinet-tiny-l. Scholars can choose the model size according to their own needs. The weight files of each model in the above image are in the trained_models folder in the source code package. You can directly use these weight files for testing. .

3.1 Training and testing commands

The following are the training and testing commands corresponding to each model:

############### Table 5, Baseline 
# train, or generate maps without conversion (uncomment the --evaluate)
python main.py --model pidinet --config baseline --sa --dil --resume --iter-size 24 -j 4 --gpu 0 --epochs 20 --lr 0.005 --lr-type multistep --lr-steps 10-16 --wd 1e-4 --savedir /path/to/table5_baseline --datadir /path/to/BSDS500 --dataset BSDS #--evaluate /path/to/table5_baseline.pth

# generate maps (no need to convert because baseline is already a vanilla cnn)
python main.py --model pidinet --config baseline --sa --dil -j 4 --gpu 0 --savedir /path/to/table5_baseline --datadir /path/to/BSDS500 --dataset BSDS --evaluate /path/to/table5_baseline.pth

# 101 FPS
python throughput.py --model pidinet --config baseline --sa --dil -j 1 --gpu 0 --datadir /path/to/BSDS500 --dataset BSDS



############### Table 5, PiDiNet
# train, or generate maps without conversion (uncomment the --evaluate)
python main.py --model pidinet --config carv4 --sa --dil --resume --iter-size 24 -j 4 --gpu 0 --epochs 20 --lr 0.005 --lr-type multistep --lr-steps 10-16 --wd 1e-4 --savedir /path/to/table5_pidinet --datadir /path/to/BSDS500 --dataset BSDS #--evaluate /path/to/table5_pidinet.pth

# generate maps with converted pidinet
python main.py --model pidinet_converted --config carv4 --sa --dil -j 4 --gpu 0 --savedir /path/to/table5_pidinet --datadir /path/to/BSDS500 --dataset BSDS --evaluate /path/to/table5_pidinet.pth --evaluate-converted

# 96 FPS
python throughput.py --model pidinet_converted --config carv4 --sa --dil -j 1 --gpu 0 --datadir /path/to/BSDS500 --dataset BSDS



############### Table 5, PiDiNet-L
# train, or generate maps without conversion (uncomment the --evaluate)
python main.py --model pidinet --config carv4 --resume --iter-size 24 -j 4 --gpu 0 --epochs 20 --lr 0.005 --lr-type multistep --lr-steps 10-16 --wd 1e-4 --savedir /path/to/table5_pidinet-l --datadir /path/to/BSDS500 --dataset BSDS #--evaluate /path/to/table5_pidinet-l.pth

# generate maps with converted pidinet
python main.py --model pidinet_converted --config carv4  -j 4 --gpu 0 --savedir /path/to/table5_pidinet-l --datadir /path/to/BSDS500 --dataset BSDS --evaluate /path/to/table5_pidinet-l.pth --evaluate-converted

# 135 FPS
python throughput.py --model pidinet_converted --config carv4 -j 1 --gpu 0 --datadir /path/to/BSDS500 --dataset BSDS



############### Table 5, PiDiNet-small
# train, or generate maps without conversion (uncomment the --evaluate)
python main.py --model pidinet_small --config carv4 --sa --dil --resume --iter-size 24 -j 4 --gpu 0 --epochs 20 --lr 0.005 --lr-type multistep --lr-steps 10-16 --wd 1e-4 --savedir /path/to/table5_pidinet-small --datadir /path/to/BSDS500 --dataset BSDS #--evaluate /path/to/table5_pidinet-small.pth

# generate maps with converted pidinet
python main.py --model pidinet_small_converted --config carv4 --sa --dil -j 4 --gpu 0 --savedir /path/to/table5_pidinet-small --datadir /path/to/BSDS500 --dataset BSDS --evaluate /path/to/table5_pidinet-small.pth --evaluate-converted

# 161 FPS
python throughput.py --model pidinet_small_converted --sa --dil --config carv4 -j 1 --gpu 0 --datadir /path/to/BSDS500 --dataset BSDS



############### Table 5, PiDiNet-small-l
# train, or generate maps without conversion (uncomment the --evaluate)
python main.py --model pidinet_small --config carv4 --resume --iter-size 24 -j 4 --gpu 0 --epochs 20 --lr 0.005 --lr-type multistep --lr-steps 10-16 --wd 1e-4 --savedir /path/to/table5_pidinet-small-l --datadir /path/to/BSDS500 --dataset BSDS #--evaluate /path/to/table5_pidinet-small-l.pth

# generate maps with converted pidinet
python main.py --model pidinet_small_converted --config carv4 -j 4 --gpu 0 --savedir /path/to/table5_pidinet-small-l --datadir /path/to/BSDS500 --dataset BSDS --evaluate /path/to/table5_pidinet-small-l.pth --evaluate-converted

# 225 FPS
python throughput.py --model pidinet_small_converted --config carv4 -j 2 --gpu 0 --datadir /path/to/BSDS500 --dataset BSDS



############### Table 5, PiDiNet-tiny
# train, or generate maps without conversion (uncomment the --evaluate)
python main.py --model pidinet_tiny --config carv4 --sa --dil --resume --iter-size 24 -j 4 --gpu 0 --epochs 20 --lr 0.005 --lr-type multistep --lr-steps 10-16 --wd 1e-4 --savedir /path/to/table5_pidinet-tiny --datadir /path/to/BSDS500 --dataset BSDS #--evaluate /path/to/table5_pidinet-tiny.pth

# generate maps with converted pidinet
python main.py --model pidinet_tiny_converted --config carv4 --sa --dil -j 4 --gpu 0 --savedir /path/to/table5_pidinet-tiny --datadir /path/to/BSDS500 --dataset BSDS --evaluate /path/to/table5_pidinet-tiny.pth --evaluate-converted

# 182 FPS 
python throughput.py --model pidinet_tiny_converted --sa --dil --config carv4 -j 2 --gpu 0 --datadir /path/to/BSDS500 --dataset BSDS



############### Table 5, PiDiNet-tiny-l
# train, or generate maps without conversion (uncomment the --evaluate)
python main.py --model pidinet_tiny --config carv4 --resume --iter-size 24 -j 4 --gpu 0 --epochs 20 --lr 0.005 --lr-type multistep --lr-steps 10-16 --wd 1e-4 --savedir /path/to/table5_pidinet-tiny-l --datadir /path/to/BSDS500 --dataset BSDS #--evaluate /path/to/table5_pidinet-tiny-l.pth

# generate maps with converted pidinet
python main.py --model pidinet_tiny_converted --config carv4 -j 4 --gpu 0 --savedir /path/to/table5_pidinet-tiny-l --datadir /path/to/BSDS500 --dataset BSDS --evaluate /path/to/table5_pidinet-tiny-l.pth --evaluate-converted

# 253 FPS
python throughput.py --model pidinet_tiny_converted --config carv4 -j 2 --gpu 0 --datadir /path/to/BSDS500 --dataset BSDS



############### Table 6, PiDiNet
# train, or generate maps without conversion (uncomment the --evaluate)
python main.py --model pidinet --config carv4 --sa --dil --resume --iter-size 24 -j 4 --gpu 0 --epochs 14 --lr 0.005 --lr-type multistep --lr-steps 8-12 --wd 1e-4 --savedir /path/to/table6_pidinet --datadir /path/to/NYUD --dataset NYUD-image --lmbda 1.3 #--evaluate /path/to/table6_pidinet.pth

# generate maps with converted pidinet
python main.py --model pidinet_converted --config carv4 --sa --dil -j 4 --gpu 0 --savedir /path/to/table6_pidinet --datadir /path/to/NYUD --dataset NYUD-image --lmbda 1.3 --evaluate /path/to/table6_pidinet.pth --evaluate-converted

# 66 FPS
python throughput.py --model pidinet_converted --sa --dil --config carv4 -j 1 --gpu 0 --datadir /path/to/NYUD --dataset NYUD-image



############### Table 7, PiDiNet
# train, or generate maps without conversion (uncomment the --evaluate)
python main.py --model pidinet --config carv4 --sa --dil --resume --iter-size 24 -j 4 --gpu 0 --epochs 14 --lr 0.005 --lr-type multistep --lr-steps 8-12 --wd 1e-4 --savedir /path/to/table7_pidinet --datadir /path/to/Multicue/multicue_v2 --dataset Multicue-boundary-1 #--evaluate /path/to/table7_pidinet.pth

# generate maps with converted pidinet
python main.py --model pidinet_converted --config carv4 --sa --dil -j 4 --gpu 0 --savedir /path/to/table7_pidinet --datadir /path/to/Multicue/multicue_v2 --dataset Multicue-boundary-1 --evaluate /path/to/table7_pidinet.pth --evaluate-converted

# 17 FPS
python throughput.py --model pidinet_converted --sa --dil --config carv4 -j 1 --gpu 0 --datadir /path/to/Multicue/multicue_v2 --dataset Multicue-boundary-1

The definition of each model is located in the code as follows:

For actual training, you only need to modify the model saving path in the training command: –savedir; the data set path: –dataset

The specific usage is as follows. Enter the command in the terminal. The following is the training process:

3.1 Model saving

The trained model will be automatically saved in the specified path:

4. Test

See 3.1 above for the test commands corresponding to each model.

In the source code package I provided, the GPU test script and the CPU test script are separately distinguished, as follows:

When using
the test command, you need to modify the path to read the test set: –datadir; the path to save the test results: –savedir; the path to the trained model weights: –evaluate. Other parameters can be adjusted according to the situation.

Examples of GPU test commands are as follows:

python main_GPU.py --model pidinet_converted --config carv4 --sa --dil -j 4 --gpu 0 --savedir ./path/to/savedir_table7_pidinet --datadir ./path/to/custom_images --dataset Custom --evaluate ./trained_models/table5_baseline.pth --evaluate-converted

Examples of CPU test commands are as follows:

python main_CPU.py --model pidinet_tiny_converted_small5 --config carv4 --sa --dil -j 4 --savedir ./results/pidinet_tiny_small5_Crop --datadir ./path/to/custom_images --dataset Custom --evaluate ./path/to/Train_Models/pidinet_tiny_small5_Crop/save_models/checkpoint_008.pth --evaluate-converted

The following is the actual test process of entering commands in the terminal:

The final test results will be automatically saved in the path specified by –savedir, as follows:

4.1 Test results

4.1.1 Test scenario 1

4.1.2 Test scenario 2

4.1.3 Test scenario 3

4.2 Detect edges and fuse with original image

The fusion method is to first overlay and fuse the original image and the edge image on the Y channel alone, and finally use the fused Y channel to merge the Cr and Cb of the original image, and then convert it into a BGR image. The image obtained by this fusion method is more natural and the color information will not be lost.

4.2.1 Fusion code

# 此方法可行，在Y通道上将图像的细节融合后再与CrCb色彩通道合并颜色

import cv2
import numpy as np

# 读取原始图像

image = cv2.imread('path/to/DenoiseImages/23.jpg', cv2.IMREAD_COLOR)    # 原图

image_edge = cv2.imread("results/Table_pidinet_tiny_l/eval_results/imgs_epoch_019/23.png")  # 边缘图像

# 转YCrCb图像
image_YCrCb = cv2.cvtColor(image,cv2.COLOR_BGR2YCrCb)
image_edge_YCrCb = cv2.cvtColor(image_edge,cv2.COLOR_BGR2YCrCb)

# 拆分YCRCb图像各个通道

image_y,image_cr,image_cb = cv2.split(image_YCrCb)

image_edge_y,image_edge_cr,image_edge_cb = cv2.split(image_edge_YCrCb)

# 将边缘图像转换为彩色图像
# edge_image_color = cv2.cvtColor(image_edge, cv2.COLOR_GRAY2BGR)

# 将边缘图像叠加到原始图像上
enhanced_image_y = cv2.addWeighted(image_y, 0.9, image_edge_y, 0.1, 0)

# 合同通道
fusion_YCrCb = cv2.merge((enhanced_image_y,image_cr,image_cb))

# YCrCb转BGR
fusion_BGR = cv2.cvtColor(fusion_YCrCb,cv2.COLOR_YCrCb2BGR)

cv2.imwrite("results/DenoiseImages+Table_pidinet_tiny_l/23.bmp",fusion_BGR)

# fusion_RGB = cv2.cvtColor(fusion_YCrCb,cv2.COLOR_YCrCb2RGB)

# 显示增强后的图像
cv2.imshow("original",image)
cv2.imshow("image_edge",image_edge)
cv2.imshow('image_y',image_y)
cv2.imshow('image_edge_y',image_edge_y)
cv2.imshow('enhanced_image_y',enhanced_image_y)
cv2.imshow('Fusion Image', fusion_BGR)
cv2.waitKey(0)
cv2.destroyAllWindows()

4.2.2 Display of fusion results

4.2.1 Scenario 1

4.2.2 Scenario 2

5. Model optimization

The official website has provided a variety of models, but they are still far from enough on the CPU, so my optimization is mainly for speed.

5.1 Modify the number of channels and void rate

My minimum number of channels is 5 and the hole rate is 2. Other values ​​can also be set, which will affect the detection effect.

5.2 Crop convolution layer

The part where I cut the network structure is shown in the red box in the figure below. The cutting code script is the pidinet_Crop.py file in the source code package.

Delete the block_3 and block_4 parts in the corresponding source code, and modify it accordingly in the forword, as follows:

This method can increase the speed to a certain extent, but the detection effect is much worse.

6. Reasoning speed

The image resolution I tested is 480*360, computer processor: 12th Gen Intel® Core™ i7-12700H 2.30 GHz.

The official website table7_pidinet model size is 2.73M. CPU inference speed: 450ms/fps, GPU inference speed: 4ms/fps.

The official website table5-pidinet-tiny-l model size is 304K. CPU inference speed: 143.6ms/fps, GPU inference speed: 2ms/fps.

Model optimization, the number of clipping channels is 5, the hole rate is 2, and the model size is: 178K. CPU inference speed: 81ms/fps.

Model optimization. The cropped part of the convolution layer is shown in the red box in the figure below. The number of channels is 5, the hole rate is 2, and the model size is: 65.06K. CPU inference speed: 72.25ms/fps.

Summarize

The above is a detailed graphic tutorial for the edge detection algorithm PidiNet network to train its own data set and optimize inference testing. The network architecture adopts lightweight module integration, which is indeed fast, but there is still room for optimization. Scholars with more in-depth research are welcome to join us. Explore.

It’s not easy to summarize. Thank you for your support!