darkface 转 yolo 格式

# author: wujiahao
# last modified: 2022/3/12
#       convert darkface to yolo format and split train val test

import enum
import os
import cv2
import shutil
import glob
import numpy as np
from sklearn.model_selection import train_test_split

data_root = '/data1/wjh/darkface' # anndir  imgdir
dst_path = '/data1/wjh/yolo/darkface'

img_list = os.listdir(os.path.join(data_root, 'image')) # 6000
trainset, valset = train_test_split(img_list, test_size=3000) # 3000 train
valset, testset = train_test_split(img_list, test_size=2000) # 2000 test, 1000 val
for type, dts in enumerate([trainset,valset, testset]):
    if  not os.path.exists(os.path.join(dst_path, ['train', 'val', 'test'][type])):
        os.mkdir(os.path.join(dst_path, ['train', 'val', 'test'][type]))
        os.mkdir(os.path.join(dst_path, ['train', 'val', 'test'][type], 'images'))
        os.mkdir(os.path.join(dst_path, ['train', 'val', 'test'][type], 'labels'))
    for imgname in dts:
        img_file_path = os.path.join(data_root, 'image', imgname)
        img = cv2.imread(img_file_path)
        h_img, w_img, c_img = img.shape
        shutil.copy(img_file_path, os.path.join(dst_path, ['train', 'val', 'test'][type], 'images', imgname))
        with open(os.path.join(data_root, 'label', imgname.split('.')[0]+'.txt'), 'r') as f:
            ann = [[int(i) for i in line.strip().split()] for line in f.readlines()[1:]] # abs l t r b, should convert to rel x_center y_center w h
            # vis
            # for line in ann:
            #     cv2.rectangle(img, (line[0], line[1]), (line[2], line[3]), color=[0,255,0], thickness=2, lineType=4)
            # cv2.imwrite('try.png', img)
            # exit(0)
        ann = [[(line[0]+line[2])/2/w_img, (line[1]+line[3])/2/h_img, (line[2]-line[0])/w_img, (line[3]-line[1])/h_img] for line in ann]
        out_lines = [f'{0} {line[0]} {line[1]} {line[2]} {line[3]}\n' for line in ann]
        with open(os.path.join(dst_path, ['train', 'val', 'test'][type], 'labels', imgname.split('.')[0]+'.txt'), 'w') as f:
            f.writelines(out_lines)