将bdd100k数据集转换为yolo格式
因为原先从COCO中提取的数据集的检测效果一直不好，在观察了几个数据集后发现，COCO数据集的图片中我需要的目标都比较小，且我需要的一些类样本不均衡，比如truck类就特别少，尽管YOLOv4的论文中说可以使用focal loss来解决样本不均衡，但是最简单的还是换一个数据集，因此打算新做一个数据集。
官网：https://bdd-data.berkeley.edu/ https://arxiv.org/abs/1805.04687 https://github.com/bdd100k/bdd100k
BDD100K是伯克利大小AI实验室发布的一个自动驾驶领域的数据集。具体内容可以看论文，在此我只关心我需要的内容，各个类别的数据量。
由于BDD100K的数据集格式和darknet需要的yolo格式不一样，因此需要将label进行转换，顺便对label进行一下提取，去掉几个我不需要的类。Scalabel Format ，此外还有一些BDD100K自己的格式，不过与我的需求没有关系。因此先来看一下scalabel格式
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 - name: string - url: string - videoName: string (optional) - attributes: a dictionary of frame attributes - intrinsics - focal: [x, y] - center: [x, y] - nearClip: - extrinsics - location - rotation - timestamp: int64 (epoch time ms) - frameIndex: int (optional, frame index in this video) - labels [ ]: - id: string - category: string (classification) - manualShape: boolean - manualAttributes: boolean - score: float - attributes: a dictionary of label attributes - box2d: - x1: float - y1: float - x2: float - y2: float - box3d: - alpha: - orientation: - location: () - dimension: (3D point, height, width, length) - poly2d: an array of objects, with the structure - vertices: [][]float (list of 2-tuples [x, y]) - types: string - closed: boolean
再看一看实际文件，还是符合的
参考了github的一个现有转换脚本 进行了一定的修改，添加了命令行支持和类别提取。代码如下。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 import osimport jsonfrom tqdm import tqdmimport cv2TRAIN_LABEL_NAME = 'bdd100k_labels_images_train.json' VAL_LABEL_NAME = 'bdd100k_labels_images_val.json' classes = [ "person" , "bike" , "car" , "motor" , "bus" , "truck" , "rider" , "train" ] counter = {} for c in classes: counter[c] = 0 def get_args (): import argparse parser = argparse.ArgumentParser() parser.add_argument('-i' , '--images' , help ="images path" , type =str , default='./images/100k' ) parser.add_argument('-l' , '--labels' , help ="labels path" , type =str , default="./labels" ) parser.add_argument('-t' , '--type' , help ="type of dataset" , choices=['train' , 'val' ], default='train' ) return parser.parse_args() def convertBdd100k2yolo (imageFileName, label ): img = cv2.imread(imageFileName) width, height = img.shape[1 ], img.shape[0 ] dw = 1.0 /width dh = 1.0 /height catName = label['category' ] classIndex = classes.index(catName) roi = label['box2d' ] w = roi['x2' ]-roi['x1' ] h = roi['y2' ]-roi['y1' ] x_center = roi['x1' ] + w/2 y_center = roi['y1' ] + h/2 x_center, y_center, w, h = x_center*dw, y_center*dh, w*dw, h*dh return "{} {} {} {} {}\n" .format (classIndex, x_center, y_center, w, h) if __name__ == '__main__' : args = get_args() if args.type == 'train' : imageRootPath = os.path.join(args.images, 'train' ) labelFilePath = os.path.join(args.labels, TRAIN_LABEL_NAME) else : imageRootPath = os.path.join(args.images, 'val' ) labelFilePath = os.path.join(args.labels, VAL_LABEL_NAME) with open (labelFilePath) as file: lines = json.load(file) print ("loaded labels" ) for line in tqdm(lines): name = line['name' ] labels = line['labels' ] imagePath = os.path.join(imageRootPath, name) labelPath = imagePath.replace('jpg' , 'txt' ) if not os.path.isfile(os.path.realpath(imagePath)): continue with open (labelPath, 'w' ) as file: for label in labels: cat = label['category' ] if cat in classes: counter[cat] += 1 file.write(convertBdd100k2yolo(imagePath,label)) print (counter)
效果如下，训练集因为中途缩放了一下vscode窗口，所以进度条就这样了。除去一个20%量的test，类别数量还是和论文里面的差不多的。
经过训练发现存在一些无标签图片，因此需要进行删除，代码如下
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 import osdef get_args (): import argparse parser = argparse.ArgumentParser() parser.add_argument('-i' , '--images' , help ="images path" , type =str , default='./images/100k' ) parser.add_argument('-t' , '--type' , help ="type of dataset" , choices=['train' , 'val' ], default='train' ) return parser.parse_args() def getFileList (dir , extract fileList = [] filenames = os.listdir(dir ) for filename in filenames: ext = os.path.splitext(filename)[-1 ] if ext == extract: fileList.append(filename) return fileList if __name__ == '__main__' : args = get_args() if args.type == 'train' : imageRootPath = os.path.join(args.images, 'train' ) else : imageRootPath = os.path.join(args.images, 'val' ) imageName = getFileList(imageRootPath, '.jpg' ) labelName = getFileList(imageRootPath, '.txt' ) labelName = [label.replace(".txt" , ".jpg" ) for label in labelName] lackImages = set (imageName) - set (labelName) for file in lackImages: os.remove(os.path.join(imageRootPath,file))
标签转换完成，接下来生成train.txt和val.txt就和之前的coco2yolo一样了。随后的效果就得训练后才能知道了。
[1]Yu F, Chen H, Wang X, et al. Bdd100k: A diverse driving dataset for heterogeneous multitask learning[C]//Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2020: 2636-2645.BDD100K：一个大规模、多样化的驾驶视频数据集 bdd100k_to_yolo_converter
