BDD100K自动驾驶数据集格式转YOLO格式

说明：

为了用BDD100K数据集训练YOLOV5模型，首先需要将BDD100K数据集格式转成YOLOV5支持的输入格式。转换代码如下：

一、BDD100K转YOLO格式

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import re
import os
import json


def search_file(data_dir, pattern=r'\.jpg$'):
    root_dir = os.path.abspath(data_dir)
    for root, dirs, files in os.walk(root_dir):
        for f in files:
            if re.search(pattern, f, re.I):
                abs_path = os.path.join(root, f)
                # print('new file %s' % absfn)
                yield abs_path


class Bdd2yolov5:
    def __init__(self):
        self.bdd100k_width = 1280
        self.bdd100k_height = 720
        # self.select_categorys = ["person", "car", "bus", "truck"]
        # self.cat2id = {
        #     "person": 0,
        #     "car": 1,
        #     "bus": 1,
        #     "truck": 1
        # }
        self.select_categorys = ["traffic sign", "tf_green", "tf_yellow", "tf_red", "tf_none"]  # 要用到的类别名称
        self.cat2id = {
            "traffic sign": 0,
            "tf_green": 1,
            "tf_yellow": 2,
            "tf_red": 3,
            "tf_none": 4
        }

    @property
    def all_categorys(self):
        return ["person", "rider", "car", "bus", "truck", "bike",
                "motor", "traffic light", "traffic sign", "train"]

    def _filter_by_attr(self, attr=None):
        if attr is None:
            return False
            # 过滤掉晚上的图片
        if attr['timeofday'] == 'night':
            return True
        return False

    def _filter_by_box(self, w, h):
        # size ratio
        # 过滤到过于小的小目标
        threshold = 0.001
        if float(w * h) / (self.bdd100k_width * self.bdd100k_height) < threshold:
            return True
        return False

    def bdd2yolov5(self, path, save_txt_path):  # 输入一张图片的标签
        lines = ""
        with open(path) as fp:
            j = json.load(fp)
            # if self._filter_by_attr(j['attributes']):  # 过滤掉晚上的图片
            #     return
            for fr in j["frames"]:
                dw = 1.0 / self.bdd100k_width
                dh = 1.0 / self.bdd100k_height
                for obj in fr["objects"]:
                    category = obj["category"]
                    if (category == "traffic light"):
                        color = obj['attributes']['trafficLightColor']
                        category = "tf_" + color

                    if category in self.select_categorys:
                        idx = self.cat2id[category]
                        cx = (obj["box2d"]["x1"] + obj["box2d"]["x2"]) / 2.0
                        cy = (obj["box2d"]["y1"] + obj["box2d"]["y2"]) / 2.0
                        w = obj["box2d"]["x2"] - obj["box2d"]["x1"]
                        h = obj["box2d"]["y2"] - obj["box2d"]["y1"]
                        if w <= 0 or h <= 0:
                            continue
                        # if self._filter_by_box(w, h): # 过滤掉过于小的小目标
                        #     continue
                        # 根据图片尺寸进行归一化
                        cx, cy, w, h = cx * dw, cy * dh, w * dw, h * dh
                        line = f"{idx} {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}\n"
                        lines += line
                if len(lines) != 0:
                    # 转换后的以*.txt结尾的标注文件放到指定目录save_txt_path位置
                    yolo_txt = path.replace(".json", ".txt")
                    yolo_txt = os.path.join(save_txt_path, yolo_txt.rsplit("/", 1)[-1])
                    with open(yolo_txt, 'w') as fp2:
                        fp2.writelines(lines)
                # print("%s has been dealt!" % path)


if __name__ == "__main__":
    bdd_label_dir = "/home/wsy/data/BDD100K/bdd100k_labels/bdd100k/labels/100k/val"
    save_txt_path = "/home/wsy/data/BDD100K/save_temp/val"  # 指定转换后生成的文件存储的目录
    cvt = Bdd2yolov5()
    for path in search_file(bdd_label_dir, r"\.json$"):
        cvt.bdd2yolov5(path, save_txt_path)

二、需要喂入YOLOV5网络训练时要指定的train.txt和val.txt生成

前提了解：
在data.yaml文件中没有定义labels的路径，那么 YOLOV5是怎么定义label路径的呢？
如果图片的路径为：
…/coco128/images/train2017/img001.jpg
那么labels必须为：
…/coco128/labels/train2017/img001.txt
一张图片对应一个txt文件，路径除了将最后一个images换成labels，没有任何不同不然就会找不到labels。

生成train.txt和val.txt，代码如下：
train.txt文件里每一行代表每张训练图片的具体路径。
val.txt文件里每一行代表每张验证图片的具体路径。

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os

if __name__ == '__main__':
    # 所有训练标签的文件夹路径
    path_label_files_train = "/home/wsy/data/BDD100K/dataSets/labels/train"
    path_label_files_val = "/home/wsy/data/BDD100K/dataSets/labels/val"
    # single_label_path = "/home/wsy/code/yolov5-6.1/BDD100_data/labels_traffic_sign_train/468724af-6af83d8d.txt"
    # 训练图像数据集和验证图像数据集的文件夹路径
    train_imageSet_file_path = "/home/wsy/data/BDD100K/dataSets/images/train"
    val_imageSet_file_path = "/home/wsy/data/BDD100K/dataSets/images/val"
    # print(labels_lists)
    # 存储训练图像、验证图像数据集路径的文件
    train_image_lists = "/home/wsy/code/yolov5-6.1/BDD100_data/ImageSets_path/train.txt"
    val_image_lists = "/home/wsy/code/yolov5-6.1/BDD100_data/ImageSets_path/val.txt"
    if not os.path.exists(train_image_lists):
        os.system(r"touch {}".format(train_image_lists))
    if not os.path.exists(val_image_lists):
        os.system(r"touch {}".format(val_image_lists))

    labels_train_lists = os.listdir(path_label_files_train)  # 读取所有标签文件名并存在一个列表当中
    labels_val_lists = os.listdir(path_label_files_val)  # 读取所有标签文件名并存在一个列表当中

    # 训练集标签集合文件生成
    with open(train_image_lists, 'w') as file_train:
        for label in labels_train_lists:
            label_name = label.rsplit("/", 1)[-1].split(".")[0] + ".jpg"
            line_name = os.path.join(train_imageSet_file_path, label_name)
            file_train.write(line_name) # 在txt文件中写入一行内容
            file_train.write("\n") # 换行

    # 验证集标签集合文件生成
    with open(val_image_lists, 'w') as file_val:
        for label2 in labels_val_lists:
            label_name2 = label2.rsplit("/", 1)[-1].split(".")[0] + ".jpg"
            line_name2 = os.path.join(val_imageSet_file_path, label_name2)
            file_val.write(line_name2) # 在txt文件中写入一行内容
            file_val.write("\n") # 换行