神经网络做轮廓识别：第二部分-模型训练与应用

本文链接：https://blog.youkuaiyun.com/ww_Hui123/article/details/137093308

本文详细介绍了如何使用神经网络进行轮廓识别，包括下载和修改PIDINET代码、创建依赖包、准备自定义数据集、模型训练以及应用模型处理新图像的过程。作者分享了个性化需求下的实践步骤和工具使用技巧。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

前言：

这篇文章紧接上一篇“神经网络做轮廓识别：第一部分-GPU部署”，上一篇我们已经完成了前置条件，现在开始训练自己的模型，并实际应用！话不多说，正文开始。

第一步：

首先要将开头提到的论文的代码下载下来，Github上不去的话也可以在我上传的资源中下载，名称为“pidinet-master”，我的这个跟原作者的程序略有不同，主要修改是取消了图片中间层的输出以及mat数据输出，另外更改了数据导入的路径。

为了确保模型所需的库正常运行，先打开cmd，切换到项目所在目录，使用以下代码打包生成依赖包：

pip freeze > requirements.txt

接下来再根据依赖包批量安装这些包，命令如下：

pip install -r requirements.txt

第二步：

上面这些步骤操作完成后就可以进行下一步，也是最麻烦的一步：准备自己的数据集。

训练模型，简单来说就是先用标准问题和答案不断让模型巩固自己的知识面，在这个模型中，问题就是原始图像，而答案就是与之相对应的轮廓图像（这里的轮廓一般是人工提取的，因此是标准答案），这些图像对称为训练集，训练的过程还需要额外准备少量标准问题和答案，用于测验或者考试，判断它学习的效果，这些数据就称为测试集。

当模型把所有的这些标准问题和答案之间的关系学习掌握之后，就算是模型训练完毕。接下来，当你给它一个新的图像时，他就会根据之前学习的内容，生成与之对应的轮廓图像了。

因此标准的图像和与之对应的轮廓图像非常关键，这些数据一部分构成训练集，一部分构成测试集。网上已经有人整理了标准的数据集，比如下面这两个数据集。里面包含大量日常的图片，以及与之相对应的人工提取的轮廓图像，可以方便你训练日常图像的轮廓提取模型，例如人体、物体等。

标准数据集

在这里我就不用这些数据集了，我会生成我自己的数据集，因为我的图像和数据集给出的图像风格差异太大，用那些数据训练出来的模型不太适用。

我的问题的背景是，用高速相机记录了水滴撞击在一个平面及其后续运动和变形的大量图像，我希望把水滴所覆盖的区域确定出来，即有水的地方赋值为1，没有水滴的地方赋值为0，最后和其他的物理量相乘，把水滴以外的数据精确的设为0值，只保留水滴所在范围内的物理参数的数据。

下面是我需要处理的原始图像：

首先生成自己的数据集，也就是手动抠图，人工把水滴轮廓确定出来。由于我的图像特征比较简单，因此不用非常大的数据量，我陆陆续续一共抠了510个图像，因此我的数据即就包含510对原始图像及其对应的轮廓图像。

在这里，我的轮廓图像是填充成白色的，不是单纯的把轮廓画出来，而是轮廓包围的范围用白色填充，如下图所示。大家可以根据自己的需求创建自己的数据集。

另外，在这里我扣图时使用了光影魔术手和按键精灵这两个软件，通过按键精灵录制脚本，用于进行图像保存、切换下一张选择抠图按钮等操作，只有具体画轮廓线时我手动操作，画完之后用脚本保存并切换下一张，这样效率大大提高！下面这些后缀带“副本”的就是抠完图保存的图像，但我没有直接使用这样的结果，而是进一步，使用用Python将抠完的图像轮廓以外赋值0（即黑色），轮廓内赋值255（即白色），实现这一效果的程序也非常简单，这才是我最终使用的图像。

人工标记完轮廓的图像用cv2.threshold与cv2.drawContours就可以方便的将其中填充成白色。我上传的资源中有一个“Make_train_dataset.py”文件，其中包含几个类似的提取轮廓的函数，可以参考，但自己使用时需要根据具体情况进行修改。

第三步：

数据集有了之后就可以着手开始训练，将自己的数据集放在data\BSDS500\HED-BSDS\train路径下，一定要在这个目录下，否则需要改源程序中的部分路径内容。

接着创建list文件，用于确定哪些图片时训练集，哪些时测试集，在这里我用Python程序生成这两个列表（如果数据存储路径跟我不一样的话需要根据实际情况修改路径），这个make_list函数输入轮廓图所在的路径，他会读取所有轮廓图的名称，根据这个名称分确定原图和轮廓图的路径，当然如果数量不多的话手打也可以：

def make_list(file_path):
    #创建list
    filelist = getFilePathList(os.path.join(file_path,'Png_fill'),'.png')
    result=[0 for t in range(len(filelist))]
    for i in range(len(filelist)):
        result[i] =  os.path.join('train\MyData\Jpg',filelist[i][:-4]+'.jpg') + ' '\
            + os.path.join('train\MyData\Png_fill',filelist[i][:-4]+'.png')
    result_train=''
    result_test=''
    I_tr=0
    I_te=0
    for i in range(len(filelist)):
        if random.random()<0.8:
            result_train=result_train+result[i]+'\n'
            I_tr=I_tr+1
        else:
            result_test=result_test+result[i]+'\n'
            I_te=I_te+1
    print("I_tr,I_te = ",I_tr,I_te ,"Ratio = %.3f" % (I_tr/len(filelist)))
    
    return result_train,result_test

随后回到上一级目录，在这里创建或编辑这两个.lst文件，将程序生成的result_train和result_test列表分别复制到这两个.lst文件内。

这里我没有它原本的list文件名，而是自己定义了新的文件名（即：0_my_train_pair.lst和0_my_test_pair.lst），如果这里自己改了名称，那么就需要打开edge_dataloader.py文件，修改下图所示部分的文件名名称，以便程序可以找到自己定义的数据集列表。

接下来就可以输入命令训练模型了！

打开cmd，切换到pidinet-master目录下，运行以下代码，每个参数的含义在main.py文件中有说明，重点关注“--savedir trained_models\my_pidinet1”，这里指定生成的模型存放的路径，然后回车运行，同时监控GPU状态。

一般第一次运行可能会有奇怪的问题，比如我就出现了半个屏幕花了的问题，关闭软件重新运行就好了。

python main.py --model pidinet --ablation --only-bsds --config carv4 --sa --dil --resume --iter-size 25 -j 4 --gpu 0 --epochs 30 --lr 0.005 --lr-type multistep --lr-steps 10-16 --wd 1e-4 --savedir trained_models\my_pidinet1 --datadir data\BSDS500\HED-BSDS --dataset BSDS

训练完成后，模型会存在trained_models路径下，到这一步就算模型训练完成。

第四步：

接下来，就轮到应用模型来处理自己的图片了！

首先确定数据的路径，取几张图片为例，图片存在data\New_data\Jpg路径下，输出结果定在Result中。

然后在命令行输入以下命令，其中--savedir就是结果保存的路径，--datadir是数据输入的路径，--evaluate trained_models\my_pidinet7\save_models\checkpoint_029.pth是选择哪个模型，我用以前训练好的my_pidinet7下的checkpoint_029.pth这个模型，自己的模型叫什么名字就输入什么名字，及其对应的路径。

python main.py --model pidinet_converted --config carv4 --sa --dil -j 4 --gpu 0 --savedir E:\Edge_Detection\pidinet-master\data\New_data\Result --datadir E:\Edge_Detection\pidinet-master\data\New_data\Jpg --dataset Custom --evaluate trained_models\my_pidinet7\save_models\checkpoint_029.pth --evaluate-converted