caffe-ssd训练自己的数据

1.配置环境

参考上一篇博客：cuda:9.0-cudnn7-devel-ubuntu16.04 python3 docker 安装 caffe

2.准备数据

2.1 获取数据

这次是要做一个车身条幅检测的项目。大部分数据从百度图片爬取，少量通过微博或者截取视频帧获取。因为违规的车身条幅数据较难获取，所以大部分图片都是一些运送救灾物资的车辆。

2.2 标注数据

使用labelImg对图片进行标注。因为检测类别只有一个，所以设置默认标签和自动保存之后还是标的比较快的。这里有两个地方要注意：

因为打开图片进行标记时默认全屏（也可能是设置的问题？）而一些小图片上面的物体本来就看不清，强行放大之后进行标记的话可能对模型的准确率有影响（没有标记difficult属性），所以我在标记的时候都是先将图片显示为原图大小。可以修改labelImg.py设置显示原图大小的快捷键。
因为标记框都是水平的矩形，而只要条幅倾斜一点，标记框内大部分区域都是没有物体的。所以后面数据增广如果有旋转操作的话，之前的标记框就不适用了，只能重新标记。这里建议先做旋转的增广，统一标注之后再做其他的增广。

2.3 数据增广

使用这个repo做数据增广。我做的增广包括小角度旋转、水平镜像和修改hsv值。根据里面的quick-start.ipynb很容易修改适来配自己的数据。最后得到训练集3041张和测试集605张。截取部分代码：

try:
    import xml.etree.cElementTree as ET
except ImportError:
    import xml.etree.ElementTree as ET
# 解析xml文件
def getAnnotBoxLoc(AnotPath):
    tree = ET.ElementTree(file=AnotPath)
    root = tree.getroot()
    ObjectSet=root.findall('object')
    ObjBndBoxSet={}
    for Object in ObjectSet:
        ObjName=Object.find('name').text
        BndBox=Object.find('bndbox')
        x1 = float(BndBox.find('xmin').text)
        y1 = float(BndBox.find('ymin').text)
        x2 = float(BndBox.find('xmax').text)
        y2 = float(BndBox.find('ymax').text)
        BndBoxLoc=[x1,y1,x2,y2]
        if ObjBndBoxSet.__contains__(ObjName):
            ObjBndBoxSet[ObjName].append(BndBoxLoc)
        else:
            ObjBndBoxSet[ObjName]=[BndBoxLoc]
    if len(ObjBndBoxSet.keys()) == 1:
        return np.array(ObjBndBoxSet[ObjName])
    else:
        return ObjBndBoxSet

data_path = ''
xml_path = ''
aug_data_path = ''
aug_xml_path = ''
train_image_list = os.listdir(data_path)
print(len(train_image_list))

for image_name in train_image_list:
    img = cv2.imread(os.path.join(data_path, image_name))
    bboxes = getAnnotBoxLoc(os.path.join(xml_path, image_name.split('.')[0]+'.xml'))
    img_, bboxes_ = RandomHorizontalFlip(1)(img.copy(), bboxes.copy())
    bboxes_.astype(np.int64)
#     保存图片
    cv2.imwrite(os.path.join(aug_data_path, image_name.split('.')[0]+'_hf.jpg'), img_)
#     保存标注
    tree = ET.ElementTree(file=os.path.join(xml_path, image_name.split('.')[0]+'.xml'))
    root = tree.getroot()
    
    name = root.find('filename')
    name.text = image_name.split('.')[0]+'_hf.jpg'
    
    objectSet = root.findall('object')
    for i,obj in enumerate(objectSet):
        bndBox = obj.find('bndbox')
        bndBox.find('xmin').text = str(int(bboxes_[i][0]))
        bndBox.find('ymin').text = str(int(bboxes_[i][1]))
        bndBox.find('xmax').text = str(int(bboxes_[i][2]))
        bndBox.find('ymax').text = str(int(bboxes_[i][3]))
    tree.write(os.path.join(aug_xml_path, image_name.split('.')[0]+'_hf.xml'))

3.处理数据

在caffe的data目录下新建banner（自己的数据集）和VOCdevkit两个目录，使data目录结构如下：

.
├── banner
├── cifar10
├── coco
├── ilsvrc12
├── ILSVRC2016
├── mnist
├── VOC0712
└── VOCdevkit
    └── banner
        ├── Annotations（所有的标注文件）
        ├── ImageSets
        │   └── Main（里面有test.txt和trainval.txt两个文件，文件内容分别是测试集和训练集全部图片的文件名，每个文件名占一行，一般由matlab脚本生成，这里是我自己分的）
        ├── JPEGImages（保存所有的图片）
        └── lmdb（最后生成的lmdb文件存放目录）
            ├── banner_test_lmdb
            └── banner_trainval_lmdb

data/banner下包含以下文件：

labelmap_voc_banner.prototxt

item {
  name: "none_of_the_above"
  label: 0
  display_name: "background"
}
item {
  name: "banner"
  label: 1
  display_name: "banner"
}

test_name_size.txt(由create_list_banner.sh生成)

test0259 360 480
test0375 1457 1080
test0118 350 500
test0056 750 1334
...

test.txt(同上)

banner/JPEGImages/test0259.jpg banner/Annotations/test0259.xml
banner/JPEGImages/test0375.jpg banner/Annotations/test0375.xml
banner/JPEGImages/test0118.jpg banner/Annotations/test0118.xml
banner/JPEGImages/test0056.jpg banner/Annotations/test0056.xml
...

trainval.txt(同上)

banner/JPEGImages/train_2878.jpg banner/Annotations/train_2878.xml
banner/JPEGImages/train_2904.jpg banner/Annotations/train_2904.xml
banner/JPEGImages/train_0786.jpg banner/Annotations/train_0786.xml
banner/JPEGImages/train_1818.jpg banner/Annotations/train_1818.xml
...

create_list_banner.sh由data/VOC0712中的create_list.sh（看得懂大概的意思，有些细节还不清楚）修改而来，修改之后mv到data/banner下。执行前ImageSets/Main中需要有test.txt和trainval.txt两个文件。执行后会生成test.txt、trainval.txt、test_name_size.txt和trainval_name_size.txt4个文件。

修改同目录下的create_data.sh得到create_data_banner.sh，生成用于训练的lmdb文件。

4.训练

下载SSD页面的预训练模型：07++12+COCO: SSD300*（07++12表示用10k的trainval2007+test2007和16k的trainval2012作为训练集（07++12），test2012作为测试集）。下载之后将VGGNet文件夹整个移到model目录下面。
修改里面的finetune_ssd_pascal.py，包括

文件路径
defaultbox宽高比（适配条幅）
mbox_source_layers的后缀（因为原来是在21类训练，自己的数据集只有两类，使用不同的名字表示不读取这些层（mbox_source_layers）的参数）
类别数和训练参数

因为原版ssd的脚本是在python2下写的，需要修改一些地方适配python3：

finetune_ssd_pascal.py: 在Python 3中，range()与xrange()合并为range()。因此修改文件中的xrange为range
finetune_ssd_pascal.py: test_iter = int(num_test_image/test_batch_size)
model_libs.py: 注释掉文件中，函数UnpackVariable定义中与assert相关的语句: assert len > 0
model_libs.py: 将文件中类似pad=int((3+(dilation-1)*2)-1)/2的语句的/2改为//2，一共3句(python中/的结果都是float，//才是int)

使用在07++12+COCO上的预训练模型能显著加快模型训练速度，比如迭代200次就能达到59.8%的mAP。最后模型的mAP为85.8%，共迭代44000次。

完整finetune_ssd_pascal.py文件

5.测试

修改caffe/examples中的ssd_detect.ipynb，在图片上标记检测结果并保存:

# 返回检测结果
def get_result(images_dir_path):
    images_list = os.listdir(images_dir_path)
    # 对每张图片
    res_dict = dict()
    for image_name in images_list:
        image = caffe.io.load_image(os.path.join(images_dir_path, image_name))
        transformed_image = transformer.preprocess('data', image)
        net.blobs['data'].data[...] = transformed_image

        # Forward pass.
        detections = net.forward()['detection_out']

        # Parse the outputs.
        det_label = detections[0,0,:,1]
        det_conf = detections[0,0,:,2]
        det_xmin = detections[0,0,:,3]
        det_ymin = detections[0,0,:,4]
        det_xmax = detections[0,0,:,5]
        det_ymax = detections[0,0,:,6]

        # Get detections with confidence higher than 0.6.
        top_indices = [i for i, conf in enumerate(det_conf) if conf >= 0.3]

        top_conf = det_conf[top_indices] 
        top_xmin = det_xmin[top_indices] * image.shape[1]
        top_ymin = det_ymin[top_indices] * image.shape[0]
        top_xmax = det_xmax[top_indices] * image.shape[1]
        top_ymax = det_ymax[top_indices] * image.shape[0]
        
        res = np.stack([top_conf, top_xmin, top_ymin, top_xmax, top_ymax], -1).reshape(len(top_conf), 5)
        res_dict[image_name] = res
        
    return res_dict

anno_data = get_result('')

import cv2
dir = ""
images_list = os.listdir(dir)
for image_name in images_list:
    res = anno_data[image_name]
    image = cv2.imread(os.path.join(dir, image_name))
    for data in res:
        pt1 = int(data[1]), int(data[2])
        pt2 = int(data[3]), int(data[4])
        pt_txt = int(data[1]), int(data[2])-6
        pt_bg_tl = int(data[1]), int(data[2])-20
        pt_bg_br = int(data[1]) + 90, int(data[2])
        # image = cv2.rectangle(image.copy(), pt1, pt2, (0,255,0), int(max(image.shape[:2])/200))
        image = cv2.rectangle(image.copy(), pt1, pt2, (0,255,0), 2)
        image = cv2.rectangle(image.copy(), pt_bg_tl, pt_bg_br, (255, 255, 255), -1)
        cv2.putText(image, 'banner: %.2f' % data[0], pt_txt, cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 0, 0), 1)
        os.chdir("")
    cv2.imwrite(image_name, image)

测试效果图：

测试效果图（GIF）：

参考

posted @ 2019-08-05 16:44 backtosouth 阅读(943) 评论(0) 编辑收藏举报

会员力量，点亮园子希望

刷新页面返回顶部

南来北往