# bdd数据集转换为yolo格式import reimport osimport jsondef 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_pathclass Bdd2yolov5:def __init__(self):self.bdd100k_width = 1280self.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,}@propertydef 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 Truereturn 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 Truereturn Falsedef bdd2yolov5(self, path):lines = ""with open(path) as fp:j = json.load(fp)if self._filter_by_attr(j['attributes']):returnfor fr in j["frames"]:dw = 1.0 / self.bdd100k_widthdh = 1.0 / self.bdd100k_heightfor obj in fr["objects"]:if obj["category"] in self.select_categorys:idx = self.cat2id[obj["category"]]cx = (obj["box2d"]["x1"] + obj["box2d"]["x2"]) / 2.0cy = (obj["box2d"]["y1"] + obj["box2d"]["y2"]) / 2.0w = obj["box2d"]["x2"] - obj["box2d"]["x1"]h = obj["box2d"]["y2"] - obj["box2d"]["y1"]if w <= 0 or h <= 0:continueif self._filter_by_box(w, h):continue# 根据图片尺寸进行归一化cx, cy, w, h = cx * dw, cy * dh, w * dw, h * dhline = 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)
