YOLO v3实战之钢筋数量AI识别（一）_钢筋识别算法python

本次的YOLO v3实战是基于DataFountain的一个比赛：智能盘点—钢筋数量AI识别，baseline model就选用上次讲解YOLO v3理论YunYang复现的YOLO v3。本次系列也和正常我们做比赛的流程一样分为两部分，这次也是第一部分将会带大家跑通baseline（比赛的话可能会对比多个，这里仅跑YOLO v3），第二部分将会分析baseline出现的问题结合赛题背景进行改进。

目录

• 题目重述
• 数据准备
• 修改相关配置路径
• 开始训练
• 测试结果

题目重述

题目背景

• 在工地现场，对于进场的钢筋车，验收人员需要对车上的钢筋进行现场人工点根，确认数量后钢筋车才能完成进场卸货。目前现场采用人工计数的方式，如图1-1中所示：
  
• 上述过程繁琐、消耗人力且速度很慢（一般一车钢筋需要半小时，一次进场盘点需数个小时）。针对上述问题，希望通过手机拍照->目标检测计数->人工修改少量误检的方式（如图1-2）智能、高效的完成此任务：
  

主要难点

1. 精度要求高。钢筋本身价格较昂贵，且在实际使用中数量很大，误检和漏检都需要人工在大量的标记点中找出，所以需要精度非常高才能保证验收人员的使用体验。需要专门针对此密集目标的检测算法进行优化，另外，还需要处理拍摄角度、光线不完全受控，钢筋存在长短不齐、可能存在遮挡等情况。
2. 钢筋尺寸不一。钢筋的直径变化范围较大（12-32中间很多种类）且截面形状不规则、颜色不一，拍摄的角度、距离也不完全受控，这也导致传统算法在实际使用的过程中效果很难稳定。
3. 边界难以区分。一辆钢筋车一次会运输很多捆钢筋（如图1-3），如果直接全部处理会存在边缘角度差、遮挡等问题效果不好，目前在用单捆处理+最后合计的流程，这样的处理过程就会需要对捆间进行分割或者对最终结果进行去重，难度较大。
   

任务

• 赛题基于广联达公司提供的钢筋进场现场的图片和标注，希望参赛者综合运用计算机视觉和机器学习/深度学习等技术，实现拍照即可完成钢筋点根任务，大幅度提升建筑行业关键物料的进场效率和盘点准确性，将建筑工人从这项极其枯燥繁重的工作中解脱出来，评估指标采用F1指标，这个我们在后面第二部分的实战系列再针对他具体优化。

数据准备

• 数据集链接：复制这段内容后打开百度网盘手机App，操作更方便哦 链接：https://pan.baidu.com/s/1V0s9oc1_FSNCKgRkyutU2w 提取码：hh4i
• 赛题方给我们提供了250张训练图片和200张测试图片，训练文件的标注是(x_min,y_min,x_max,y_max)：
  
• 这跟我们YOLO v3训练要的标注文件有点差别，YOLO v3的标注文件的格式是(x_min,y_min,x_max,y_max,class_id)，而且要求同张图片的标注放在同一行:
  
• YunYang给我们提供了VOC版的数据集转YOLO v3标注的脚本voc_annotation.py,所以我们就先把标注文件转换成VOC格式，再运行脚本就行了。
• 我们先在训练集上划分出训练集和验证集，大概9比1，然后给训练集和测试集分别建立一个VOC格式的文件夹（ImageSets里面还有一个Main文件夹）：
  
• 我们先把各自的图片放进JPEGImage里面，然后用代码生成ImageSets的文件：

import os
train_file=open('data/train_data_VOC/ImageSets/Main/train.txt','w')
test_file=open('data/test_VOC/ImageSets/Main/test.txt','w')
for _,_,train_files in os.walk('data/train_data_VOC/JPEGImages'):
    continue
for _,_,test_files in os.walk('data/test_VOC/JPEGImages'):
    continue
for file in train_files:
    train_file.write(file.split('.')[0]+'\n')

for file in test_files:
    test_file.write(file.split('.')[0]+'\n')
1
2
3
4
5
6
7
8
9
10
11
12

• 生成的文件里面即是图片的名称。然后我们再分别运行两个脚本，先把给的标注CSV文件转换成txt，再转换成Annotations文件夹内的xml文件（第二个脚本转换的时候顺便把类别改为rebar(钢筋)）：

import csv
import os, sys
from glob import glob
from PIL import Image

src_img_dir = r'data/train_data_VOC/JPEGImages'#图片地址
src_txt_dir = r'data/yolo'#生成txt地址

img_lists = glob(src_img_dir + '/*jpg')
img_basenames = []
for item in img_lists:
    img_basenames.append(os.path.basename(item))

img_names = []
for item in img_basenames:
    temp1, temp2 = os.path.splitext(item)
    img_names.append(temp1)

c = []
filename = r'/home/cristianoc/tensorflow-yolov3/data/yolo/train.csv'
with open(filename) as f:
    reader = csv.reader(f)
    head_now = next(reader)
    l = []
    b = []
    for cow in reader:
        label = cow[0]
        l.append(label)
        bbox = cow[1]
        b.append(bbox)
label = []
for item in l:
    temp1, temp2 = os.path.splitext(item)
    label.append(temp1)

for img in img_names:
    img_file = src_txt_dir + os.sep + img + '.txt'
    fp = open(img_file, 'w')
    for i in range(len(label)):
        if label[i] == img:
            fp.write(str(b[i]))
            fp.write('\n')
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42

import csv
import os, sys
from glob import glob
from PIL import Image

src_img_dir = r'data/train_data_VOC/JPEGImages'
src_txt_dir = r'data/yolo'
src_xml_dir = r'data/train_data_VOC/Annotations'

img_lists = glob(src_img_dir + '/*jpg')
img_basenames = []
for item in img_lists:
    img_basenames.append(os.path.basename(item))

img_names = []
for item in img_basenames:
    temp1, temp2 = os.path.splitext(item)
    img_names.append(temp1)

for img in img_names:
    im = Image.open((src_img_dir + os.sep + img + '.jpg'))
    width, height = im.size
    gt = open(src_txt_dir + os.sep + img + '.txt').read().splitlines()
    xml_file = open((src_xml_dir + os.sep + img + '.xml'), 'w')
    xml_file.write('<annotation>\n')
    xml_file.write('    <folder>VOC2007</folder>\n')
    xml_file.write('    <filename>' + str(img) + '.jpg' + '</filename>\n')
    xml_file.write('    <size>\n')
    xml_file.write('        <width>' + str(width) + '</width>\n')
    xml_file.write('        <height>' + str(height) + '</height>\n')
    xml_file.write('        <depth>3</depth>\n')
    xml_file.write('    </size>\n')

    for img_each_label in gt:
        spt = img_each_label.split(' ')  
        xml_file.write('    <object>\n')
        xml_file.write('        <name>' + str('rebar') + '</name>\n')
        xml_file.write('        <pose>Unspecified</pose>\n')
        xml_file.write('        <truncated>0</truncated>\n')
        xml_file.write('        <difficult>0</difficult>\n')
        xml_file.write('        <bndbox>\n')
        xml_file.write('            <xmin>' + str(spt[0]) + '</xmin>\n')
        xml_file.write('            <ymin>' + str(spt[1]) + '</ymin>\n')
        xml_file.write('            <xmax>' + str(spt[2]) + '</xmax>\n')
        xml_file.write('            <ymax>' + str(spt[3]) + '</ymax>\n')
        xml_file.write('        </bndbox>\n')
        xml_file.write('    </object>\n')

    xml_file.write('</annotation>')
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49

• 这样我们的VOC数据集就准备好了,然后运行脚本python scripts/voc_annotation.py --data_path data/test_VOC分别生成我们的训练标注文件和验证标注文件，这样我们的数据就准备好了。

修改相关配置路径

修改class.names文件以及__C.YOLO.CLASSES参数路径

• 我们在class文件下新建一个class.name的文件，并写入我们这次数据的唯一一个类别rebar:
  
• 然后在config.py中修改我们读入的类别路径：__C.YOLO.CLASSES = "./data/classes/class.names"

修改__C.TRAIN.ANNOT_PATH和__C.TEST.ANNOT_PATH参数路径

• 同样在config.py下修改对应训练标注文件和验证标注文件的路径，改为刚才生成好的即可。

修改__C.YOLO.ORIGINAL_WEIGHT和__C.YOLO.DEMO_WEIGHT

• __C.YOLO.ORIGINAL_WEIGHT是convert_weights.py转换权重文件的源文件，__C.YOLO.DEMO_WEIGHT是转换后生成的目标权重文件（用于将在COCO预训练好的权重文件转换后生成预训练模型）__C.YOLO.DEMO_WEIGHT就是预训练模型。
• 所以我们下载好coco权重后就执行脚本python convert_weight.py --train_from_coco开始转换：
  

开始训练

• 由于我们这次只是简单跑跑我们的baseline，所以参数我就先没调（除了调大第一阶段的学习率），一共训练50个epoch，然后用fine-tune，Warmup学习率的基本操作，这些就不讲了：
  
• 这次是云服务器训练，1080Ti真香，一下子就跑完了，我们看看结果：
  
• total_loss最后只降到了100左右，主要原因是我们的giou_loss不收敛，这个的话我们后面可以针对这个优化一下我们的算法，这个就放在本次实战的第二部分了，我们再来看看测试图片的效果。

测试结果

• 我们看看效果图和原图的对比：
  
• 感觉除了少部分的漏检之外YOLO v3的检测效果还是很不错的，这也证实了作者想在v3版本加强对小物体检测效果的思路。不过这个版本我们还没针对我们的场景进行具体的优化，具体的各种优化办法将会在下一部分讲解，希望弄懂YOLO v3理论的读者可以通过这次的实战系列真正上手YOLO v3。