caffe-ssd训练kitti、lisa数据集

来源:互联网 发布:大数据公司如何盈利 编辑:程序博客网 时间:2024/06/06 16:44

目的:将kitti、Lisa数据集合并,进行训练

一、数据集准备,将两种数据集准备成VOC格式

kitti数据集(车辆行人等):http://www.cvlibs.net/datasets/kitti/eval_object.php

lisa数据集(47种交通标志):http://cvrr.ucsd.edu/LISA/lisa-traffic-sign-dataset.html

1、准备kitti数据集

创建VOCdevkit/traffic,并为其创建子目录:Annotations,ImageSets,JPEGImages

根据VOC格式要求,创建文件夹,解压kitti数据集得到:image_2,label_2两个文件,分别为图片、标注。

标注信息保存在txt文件,内容格式:Car 0.96 0 -0.86 0.00 199.21 302.21 369.00 1.50 1.78 3.69 -3.17 1.66 3.35 -1.57

将image_2内容复制到JPEGImages目录,

在traffic目录下创建create_kitti_xml.py脚本,生成xml格式的标注,参考:https://github.com/manutdzou/KITTI_SSD/blob/master/data/KITTI/KITTI_xml.py
from xml.dom.minidom import Document
import cv2
import os

def generate_xml(name,split_lines,img_size,class_ind):
    doc = Document()  # 创建DOM文档对象

    annotation = doc.createElement('annotation')
    doc.appendChild(annotation)

    title = doc.createElement('folder')
    title_text = doc.createTextNode('KITTI')
    title.appendChild(title_text)
    annotation.appendChild(title)

    img_name=name+'.png'

    title = doc.createElement('filename')
    title_text = doc.createTextNode(img_name)
    title.appendChild(title_text)
    annotation.appendChild(title)

    source = doc.createElement('source')
    annotation.appendChild(source)

    title = doc.createElement('database')
    title_text = doc.createTextNode('The KITTI Database')
    title.appendChild(title_text)
    source.appendChild(title)

    title = doc.createElement('annotation')
    title_text = doc.createTextNode('KITTI')
    title.appendChild(title_text)
    source.appendChild(title)

    size = doc.createElement('size')
    annotation.appendChild(size)

    title = doc.createElement('width')
    title_text = doc.createTextNode(str(img_size[1]))
    title.appendChild(title_text)
    size.appendChild(title)

    title = doc.createElement('height')
    title_text = doc.createTextNode(str(img_size[0]))
    title.appendChild(title_text)
    size.appendChild(title)

    title = doc.createElement('depth')
    title_text = doc.createTextNode(str(img_size[2]))
    title.appendChild(title_text)
    size.appendChild(title)

    for split_line in split_lines:
        line=split_line.strip().split()
        if line[0] in class_ind:
            object = doc.createElement('object')
            annotation.appendChild(object)

            title = doc.createElement('name')
            title_text = doc.createTextNode(line[0])
            title.appendChild(title_text)
            object.appendChild(title)

            bndbox = doc.createElement('bndbox')
            object.appendChild(bndbox)
            title = doc.createElement('xmin')
            title_text = doc.createTextNode(str(int(float(line[4]))))
            title.appendChild(title_text)
            bndbox.appendChild(title)
            title = doc.createElement('ymin')
            title_text = doc.createTextNode(str(int(float(line[5]))))
            title.appendChild(title_text)
            bndbox.appendChild(title)
            title = doc.createElement('xmax')
            title_text = doc.createTextNode(str(int(float(line[6]))))
            title.appendChild(title_text)
            bndbox.appendChild(title)
            title = doc.createElement('ymax')
            title_text = doc.createTextNode(str(int(float(line[7]))))
            title.appendChild(title_text)
            bndbox.appendChild(title)

    # 将DOM对象doc写入文件
    f = open('./Annotations/'+name+'.xml','w')
    f.write(doc.toprettyxml(indent = ''))
    f.close()

if __name__ == '__main__':
    class_ind=('Pedestrian', 'Car', 'Cyclist')
    cur_dir=os.getcwd()
    labels_dir=os.path.join(cur_dir,'label_2')
    for parent, dirnames, filenames in os.walk(labels_dir): # 分别得到根目录,子目录和根目录下文件   
        for file_name in filenames:
            full_path=os.path.join(parent, file_name) # 获取文件全路径
            f=open(full_path)
            split_lines = f.readlines()
            name= file_name[:-4] # 后四位是扩展名.txt,只取前面的文件名
            img_name=name+'.png'
            img_path=os.path.join('/data/wuyan/data/VOCdevkit/traffic/JPEGImages',img_name) # 路径需要自行修改            
            img_size=cv2.imread(img_path).shape
            generate_xml(name,split_lines,img_size,class_ind)

执行create_kitti_xml.py,会在Annotations目录下生成xml格式的标注信息。

2、准备Lisa数据集

创建一个目录signdatabase,将Lisa数据集解压到该目录,解压后得到 多个目录(图片),其中有一个allAnnotations.csv标注文件,其内容格式为:

Filename;Annotation tag;Upper left corner X;Upper left corner Y;Lower right corner X;Lower right corner Y;Occluded,On another road;Origin file;Origin frame number;Origin track;Origin track frame number

为解析方便,删除该文件第一行内容

同样在traffic目录下编写一个脚本create_lisa_xml.py,来生成xml格式的标注文件

from xml.dom.minidom import Document
import xml.dom.minidom
import cv2
import os

def generate_xml(name,data,img_size):
    file = './Annotations/' + name + '.xml'

    #如果已经存在,读出xml内容,插入一个节点

    if os.path.exists(file):
        old_xml = xml.dom.minidom.parse(file)

        new_object = old_xml.createElement('object')
        title = old_xml.createElement('name')
        title_text = old_xml.createTextNode(data[1])
        title.appendChild(title_text)
        new_object.appendChild(title)

        bndbox = old_xml.createElement('bndbox')
        new_object.appendChild(bndbox)
        title = old_xml.createElement('xmin')
        title_text = old_xml.createTextNode(str(int(float(data[2]))))
        title.appendChild(title_text)
        bndbox.appendChild(title)
        title = old_xml.createElement('ymin')
        title_text = old_xml.createTextNode(str(int(float(data[3]))))
        title.appendChild(title_text)
        bndbox.appendChild(title)
        title = old_xml.createElement('xmax')
        title_text = old_xml.createTextNode(str(int(float(data[4]))))
        title.appendChild(title_text)
        bndbox.appendChild(title)
        title = old_xml.createElement('ymax')
        title_text = old_xml.createTextNode(str(int(float(data[5]))))
        title.appendChild(title_text)
        bndbox.appendChild(title)

        annotation = old_xml.getElementsByTagName('annotation')[0]
        annotation.appendChild(new_object)
        # 将DOM对象doc写入文件
        os.remove(file)
        f = open(file,'w')
        f.write(old_xml.toprettyxml(indent = ''))
        f.close()
    #不存在
    else:
        doc = Document()  # 创建DOM文档对象

        annotation = doc.createElement('annotation')
        doc.appendChild(annotation)

        title = doc.createElement('folder')
        title_text = doc.createTextNode('KITTI')
        title.appendChild(title_text)
        annotation.appendChild(title)

        img_name = name + '.png'

        title = doc.createElement('filename')
        title_text = doc.createTextNode(img_name)
        title.appendChild(title_text)
        annotation.appendChild(title)

        source = doc.createElement('source')
        annotation.appendChild(source)

        title = doc.createElement('database')
        title_text = doc.createTextNode('The KITTI Database')
        title.appendChild(title_text)
        source.appendChild(title)

        title = doc.createElement('annotation')
        title_text = doc.createTextNode('KITTI')
        title.appendChild(title_text)
        source.appendChild(title)

        size = doc.createElement('size')
        annotation.appendChild(size)

        title = doc.createElement('width')
        title_text = doc.createTextNode(str(img_size[1]))
        title.appendChild(title_text)
        size.appendChild(title)

        title = doc.createElement('height')
        title_text = doc.createTextNode(str(img_size[0]))
        title.appendChild(title_text)
        size.appendChild(title)

        title = doc.createElement('depth')
        title_text = doc.createTextNode(str(img_size[2]))
        title.appendChild(title_text)
        size.appendChild(title)

        object = doc.createElement('object')
        annotation.appendChild(object)
        title = doc.createElement('name')
        title_text = doc.createTextNode(data[1])
        title.appendChild(title_text)
        object.appendChild(title)

        bndbox = doc.createElement('bndbox')
        object.appendChild(bndbox)
        title = doc.createElement('xmin')
        title_text = doc.createTextNode(str(int(float(data[2]))))
        title.appendChild(title_text)
        bndbox.appendChild(title)
        title = doc.createElement('ymin')
        title_text = doc.createTextNode(str(int(float(data[3]))))
        title.appendChild(title_text)
        bndbox.appendChild(title)
        title = doc.createElement('xmax')
        title_text = doc.createTextNode(str(int(float(data[4]))))
        title.appendChild(title_text)
        bndbox.appendChild(title)
        title = doc.createElement('ymax')
        title_text = doc.createTextNode(str(int(float(data[5]))))
        title.appendChild(title_text)
        bndbox.appendChild(title)

        # 将DOM对象doc写入文件
        f = open(file,'w')
        f.write(doc.toprettyxml(indent = ''))
        f.close()

def find_last(string, str):
    last_position = -1
    while True:
        position = string.find(str, last_position + 1)
        if position == -1:
            return last_position
        last_position = position



if __name__ == '__main__':

    #只有这些类别可用

    class_ind=('addedLane', 'curveLeft', 'curveRight', 'keepRight', 'laneEnds', 'merge', 'noLeftTurn', 'noRightTurn', 'pedestrianCrossing', 'roundabout', 'signalAhead', 'stopAhead', 'turnLeft', 'turnRight', 'yield')

    with open('allAnnotations.csv', 'r') as r_tdf:
        for each_line in r_tdf:
            data = each_line.strip().split(';')
            file_path = data[0] #img file obsolute path
            tag = data[1]
            img_size = cv2.imread(file_path).shape
            name = file_path[find_last(file_path, '/') + 1:-4]
            if tag in class_ind:
                print('name : ' + name)
                generate_xml(name, data, img_size)

执行该脚本后,会在Annotations目录下生成Lisa数据集的标注文件。

最后需要将Lisa所有标注的图片复制到JPEGImages下。


数据集准备先到这里。。。下篇再写如何训练

0 0