基于红外数据集（FLIR）训练YOLOv4

1、引言

毕业设计是基于热红外图像的目标检测系统设计，老板要求做一些横向和纵向的实验对比作为支撑，自己基于项目构建的车载红外数据集只有4000多张，加上是自己手动标注的，直接用于实验对比说服力不够，加上标注的精度跟开源的标准数据集相比有一定差距，所以前期的实验还是要在公共数据集上做，最后再迁移到自己的数据集上来。

实验室之前采购了FLIR的红外摄像头，后面就关注到他们开源了一个比较大的带真实标签的红外数据集，官网地址为FLIR数据集，跟我手头数据集的应用场景非常相似，所以决定以该数据集作为前期实验的主要基准。后续的系统是要移植到NVIDIA Jetson NX上做实时目标检测，对目标检测网络的速度有比较高的要求，网上检索到目前SOTA的网络中YOLOv4各方面性能都不错，而且官网配备了详细的训练、调参、测试等教程，简直不要太贴心，所以先从YOLOv4开始我的对比之旅吧！

PS:后续如果实验顺利且有时间的话，我争取把后续的嵌入式算法移植部分也分享到网上，嘿嘿嘿

2、一个检索各领域SOTA文章的网站分享

在这个网站上可以找到各个领域目前排名靠前的文章以及其相应代码，感觉非常有用，点击这里

3、FLIR数据集

1）FLIR数据集下载地址，我在官网下载了很久都没成功，后面还是一个中科大的同学分享给我的（跪谢），有需要用于科研的同学下载不了的话可以通过博客联系我分享。

2）FLIR数据集简介：
一共有三个数据集，只有第一个缩减版的可供下载，其他两个是需要钱钱的，不过第一个对于我来说也够用了，训练集有8862张图像、测试集1366张图像，外加一个视频切割成的图像4224张。

可供下载的这个数据集只标注了person、car、bicycle、dog、other vehicle这五类目标，具体情况如下：

数据下载解压后会得到一个文件夹，里面包含了以下文件夹：

两个PDF文件是关于数据集的详细解释，我就不赘述啦，大家自己看哈

3）FLIR数据集处理
官网标签的数据格式是基于coco数据集的，也就是保存在.json文件中，要在YOLOv4上进行训练需要先转换标签的格式，我的操作步骤为先转voc格式（xml），再转YOLO格式，这部分话不多说，直接上代码，具体的解释我都放到代码里了：

json-to-xml.py

#-*-coding:utf-8-*-
import glob
import json
from lxml.etree import Element, SubElement, tostring
from xml.dom.minidom import parseString

classes = ['people', 'bicycle', 'car', 'vehicle']

# 定义json文件操作函数,读取json文件中的包围框信息和类别信息，并用于构建xml文件，这部分看不懂的可以先了解下json文件
def list_json(json_path, output_xml_path):
    json_files = glob.glob(json_path + "*.json")   # 列出文件夹中的所有后缀为.json的文件，用于批量操作

    # 对每一个.json文件进行操作
    for json_file in json_files:
        # 以只读方式打开json文件
        with open(json_file, 'r') as f:
            json_message = json.load(f)       # 解析json文件, 官网的json文件中由'annotation'和'image'组成
            annotation = json_message['annotation']
            image = json_message['image']

            objects = []                      # 用于存储'annotation'中的bbox和category
            for obj in annotation:
                # 提取每个候选框中目标的类别,我只提取了people、bicycle、car和other vehicle，因为dog的训练标签太少了
                # 类别转换基于官网的FLIR Thermal Starter Dataset Introduction文件
                category = int(obj['category_id'])
                if category == 1:
                    class_name = 'people'
                elif category == 2:
                    class_name = 'bicycle'
                elif category == 3:
                    class_name = 'car'
                elif category == 91:
                    class_name = 'vehicle'
                else:
                    continue

                # 提取每个候选框的坐标
                bbox = obj['bbox']
                # coco数据集格式为(x, y, w, h), voc数据集格式为(xmin, ymin, xmax, ymax)
                # 官网标签是依据Coco数据集形式的, 此处转换为voc格式
                xmin = int(bbox[0])             # xmin=x
                ymin = int(bbox[1])             # ymin=y
                xmax = int(bbox[0] + bbox[2])   # xmax=xmin+w
                ymax = int(bbox[1] + bbox[3])   # ymax=ymin+h
                objects.append([class_name, xmin, ymin, xmax, ymax])

            # 每一个json文件对应一张图片，前缀名一样，所以可以通过json文件的前缀来获取图片名称
            image_name = json_file[:-5].split('/')[-1]
            image_width = image['width']        # 获取图片的宽
            image_height = image['height']      # 获取图片的高

        # 生成xml文件
        write_to_xml(objects, image_name, image_width, image_height, output_xml_path)

# 构建xml文件的内容并写入相应的.xml文件中,此函数的解释参考我的另一篇博客https://blog.csdn.net/Snow_cat123456/article/details/102614507
def write_to_xml(objects, image_name, image_width, image_height, output_path):
    node_root = Element('annotation')

    node_folder = SubElement(node_root, 'folder')
    node_folder.text = 'FlirDataset'

    node_filename = SubElement(node_root, 'filename')
    node_filename.text = image_name + '.jpeg'

    node_source = SubElement(node_root, 'source')
    node_database = SubElement(node_source, 'database')
    node_database.text = 'FLIR'

    node_size = SubElement(node_root, 'size')
    node_width = SubElement(node_size, 'width')
    node_width.text = str(image_width)

    node_height = SubElement(node_size, 'height')
    node_height.text = str(image_height)

    node_depth = SubElement(node_size, 'depth')
    node_depth.text = '1'

    node_segmented = SubElement(node_root, 'segmented')
    node_segmented.text = '0'


    for obj in objects:
        node_object = SubElement(node_root, 'object')
        node_name = SubElement(node_object, 'name')
        node_name.text = obj[0]

        node_pose = SubElement(node_object, 'pose')
        node_pose.text = 'Unspecified'
        node_truncated = SubElement(node_object, 'truncated')
        node_truncated.text = '0'
        node_difficult = SubElement(node_object, 'difficult')
        node_difficult.text = '0'

        # 重点在此处也就是包围框的位置
        node_bndbox = SubElement(node_object, 'bndbox')
        node_xmin = SubElement(node_bndbox, 'xmin')
        node_xmin.text = str(obj[1])
        node_ymin = SubElement(node_bndbox, 'ymin')
        node_ymin.text = str(obj[2])
        node_xmax = SubElement(node_bndbox, 'xmax')
        node_xmax.text = str(obj[3])
        node_ymax = SubElement(node_bndbox, 'ymax')
        node_ymax.text = str(obj[4])

    xml = tostring(node_root, pretty_print=True)  # 格式化显示，该换行的换行
    dom = parseString(xml)
    savingtxtmsg = open(output_path + '/' + image_name + '.xml', 'wb')
    savingtxtmsg.write(xml)
    savingtxtmsg.close()
    print('Finished convert the ', image_name)


json_path = '/the/path/of/your/json/file/'
xml_path = '/the/path/you/want/to/save/the/xml'
list_json(json_path, xml_path)

training、validation、video三个文件夹是一样的操作，然后再将xml文件转换为yolo格式的，代码如下：

# xml_to_yolo_txt.py
# 此代码和VOC_KITTI文件夹同目录
import os
import xml.etree.ElementTree as ET
# 这里的类名为我们xml里面的类名，顺序现在不需要考虑
class_names = ['people', 'bicycle', 'car', 'vehicle']
# xml文件路径
path = '/the/path/of/the/xml/path/' 
# 转换一个xml文件为txt
def single_xml_to_txt(xml_file):
    tree = ET.parse(os.path.join(path, xml_file))
    root = tree.getroot()
    # 保存的txt文件路径
    txt_file = os.path.join('/the/path/you/want/to/save/yolo_label/', xml_file.split('.')[0]+'.txt')
    with open(txt_file, 'w') as txt_file:
        for member in root.findall('object'):
            #filename = root.find('filename').text
            picture_width = int(root.find('size')[0].text)
            picture_height = int(root.find('size')[1].text)
            class_name = member[0].text
            # 类名对应的index
            class_num = class_names.index(class_name)

            box_x_min = int(member[4][0].text) # 左上角横坐标
            box_y_min = int(member[4][1].text) # 左上角纵坐标
            box_x_max = int(member[4][2].text) # 右下角横坐标
            box_y_max = int(member[4][3].text) # 右下角纵坐标
            # 转成相对位置和宽高
            x_center = float(box_x_min + box_x_max) / (2 * picture_width)
            y_center = float(box_y_min + box_y_max) / (2 * picture_height)
            width = float(box_x_max - box_x_min) /  picture_width
            height = float(box_y_max - box_y_min) /  picture_height
            # print(class_num, x_center, y_center, width, height)
            txt_file.write(str(class_num) + ' ' + str(x_center) + ' ' + str(y_center) + ' ' + str(width) + ' ' + str(height) + '\n')
# 转换文件夹下的所有xml文件为txt
def dir_xml_to_txt(path):
    files = os.listdir(path)
    for xml_file in files:
        single_xml_to_txt(xml_file)
dir_xml_to_txt(path)

4、训练YOLOv4：
官网clone代码，yolov4，依据官网教程构建flir.data、flir.names，最重要的是flir_train.txt，这部分代码我也分享下

# for_train_val_txt.py
# 此代码和kitti_data文件夹同目录
import glob
train_path = '/训练图片所在文件夹/train_data/'
test_path = '/测试图片所在文件夹/test_data/'
# video_path = 'video_data/'
def generate_train_and_val(image_path, txt_file):
    with open(txt_file, 'w') as tf:
        for jpg_file in glob.glob(image_path + '*.jpeg'):
            tf.write(jpg_file + '\n')
generate_train_and_val(train_path, 'cfg_flir/flir_train.txt') # 生成的train.txt文件所在路径
generate_train_and_val(test_path, 'cfg_flir/flir_test.txt') # 生成的val.txt文件所在路径

修改yolov4.cfg中的classes=4、filter=27等信息，channels=1（FLIR数据集中的图片为单通道图片），然后运行命令：

./darknet detector train cfg_flir/flir.data cfg_flir/yolo-flir.cfg yolov4.conv.137

跑起来了

后续试验有进展了我再来更新