bdd100k数据集标签转换为yolo格式的txt文件

# bdd数据集转换为yolo格式
import re
import os
import json

def search_file(data_dir, pattern=r'\.jpg$'):
    # os.path.abspath(path)	返回绝对路径
    root_dir = os.path.abspath(data_dir)
    # os.walk() 目录遍历器，
    for root, dirs, files in os.walk(root_dir):
        for f in files:
            # re.search会匹配整个字符串，并返回第一个成功的匹配。
            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", "rider", "car", "bus", "truck", "bike",
                "motor"]
        self.cat2id = {
            "person": 0,
            "rider": 0,
            "car": 1,
            "bus": 3,
            "truck":3,
            "bike": 0,
            "motor": 0,
        }

    @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):
        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"]:
                    if obj["category"] in self.select_categorys:
                        idx = self.cat2id[obj["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结尾的标注文件我就直接和*.json放一具目录了
                    # yolov5中用到的时候稍微挪一下就行了
                    yolo_txt = path.replace(".json", ".txt")
                    with open(yolo_txt, 'w') as fp2:
                        fp2.writelines(lines)
                # print("%s has been dealt!" % path)


if __name__ == "__main__":
    # 输入：bdd标签位置
    bdd_label_dir = "E:/dataset3.17/bdd100k_labels/val"
    cvt = Bdd2yolov5()
    for path in search_file(bdd_label_dir, r"\.json$"):
        cvt.bdd2yolov5(path)