【数据集】LaRa交通灯数据集解析

Lara转yolov5可用的数据集；

code： Lara2coco128；

# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
# Lara dataset http://www.lara.prd.fr/benchmarks/trafficlightsrecognition
# Example usage: python train.py --data VisDrone.yaml
# parent
# ├── yolov5
# └── datasets
#     └── VisDrone  ← downloads here (2.3 GB)


# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
# path: /media/uisee/mobile_disk/workspace/ZRJ/dataset/traffic_light/Lara_UrbanSeq1_JPG # dataset root dir
# train: images  # train images (relative to 'path')   images
# val: images  # val images (relative to 'path')   images
# test: images  # test images (optional)   images
# 
# # Classes
# nc: 4  # number of classes
# names: ['go', 'stop', 'warning', 'ambiguous']
# # names: ['green', 'red', 'orange', 'ambiguous']

#Lara File format is as follows:
#Timestamp / frameindex x1 y1 x2 y2 id 'type' 'subtype'

import os
# from pathlib import Path
import shutil

def convert_box(size, box):
    # Convert Lara box to YOLO xywh box
    dw = 1. / size[0]
    dh = 1. / size[1]
    xc = ((box[0] + box[2])/2)*dw
    yc = ((box[1] + box[3])/2)*dh
    w = (box[2]-box[0])*dw
    h = (box[3]-box[1])*dh
    return xc, yc, w, h

def lara2yolo(path):
    # Path(dir / 'labels').rmdir()
    # Path(dir / 'images').rmdir()
    # (dir / 'labels').mkdir(parents=True, exist_ok=True)  # make labels directory
    # (dir / 'images').mkdir(parents=True, exist_ok=True)  # make images directory
    shutil.rmtree(os.path.join(path, 'labels'), ignore_errors=True)
    shutil.rmtree(os.path.join(path, 'images'), ignore_errors=True)
    os.mkdir(os.path.join(path, 'labels'))
    os.mkdir(os.path.join(path, 'images'))

    gt = open('Lara_UrbanSeq1_GroundTruth_GT.txt', 'rt')
    lines = gt.readlines()
    img_size = (640, 480) # (w, h)
    for i in range(13, len(lines)): 
        line = lines[i].split()
        index = line[2]
        imgname = 'frame_' + index.zfill(6) + '.jpg'
        # cls = int(line[7]) # id
        subcls = line[10]
        global cls
        if subcls == 'go':
            cls = 0
        elif subcls == 'stop':
            cls =1 
        elif subcls == 'warning':
            cls =2 
        elif subcls == 'ambiiguous':
            cls =3 
        else:
            cls =3 
        box = convert_box(img_size, tuple(map(int, line[3:7])))
        info = f"{cls} {' '.join(f'{x:.6f}' for x in box)}\n"
        labelname = './labels/' + imgname.replace('jpg', 'txt')
        imagepath = './Lara_UrbanSeq1_JPG/Lara3D_UrbanSeq1_JPG/' + imgname
        # copy images
        shutil.copy(imagepath, './images/')
        labelfile = open(labelname, 'a+')
        labelfile.write(info)
        labelfile.close()

if __name__ == "__main__":
    # Convert
    path = './dataset/traffic_light/Lara_UrbanSeq1_JPG'
    # dir = Path(path)
    lara2yolo(path)  # convert Lara_Urban annotations to YOLO labels

View Code

参考

1. ​​Lara_dataset​​;

2. ​​【Datasets】LaRa交通灯数据集解析​​;

3. ​​github_yolov5​​;

完