YOLO v3实战之钢筋数量AI识别(一)

2020-06-02 11:28:16 浏览数 (1)

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

目录

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

题目重述

题目背景

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

主要难点

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

任务

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

数据准备

  • 数据集链接: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的文件:
代码语言:javascript复制
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'):
    continuefor _,_,test_files in os.walk('data/test_VOC/JPEGImages'):
    continuefor file in train_files:
    train_file.write(file.split('.')[0] 'n')

for file in test_files:
    test_file.write(file.split('.')[0] 'n')
  • 生成的文件里面即是图片的名称。然后我们再分别运行两个脚本,先把给的标注CSV文件转换成txt,再转换成Annotations文件夹内的xml文件(第二个脚本转换的时候顺便把类别改为rebar(钢筋)):
代码语言:javascript复制
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')
代码语言:javascript复制
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>')
  • 这样我们的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_WEIGHTconvert_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。
  • ps:YOLO之父今天宣布退出计算机视觉领域,不慎感慨,本来还是十分期待YOLO v4的升级的==

0 人点赞