检查yolo格式数据集|查看删除空标签|将标签绘制在图像上并保存

# coding=utf-8
import numpy as np
import cv2
import os
import argparse
import shutil
EMPTY_TXT_NUMBER=0
def remove_empty_images_label(image_path,label_path,opt):
    if opt.copy_empty:
        os.makedirs('./empy/empyimg', exist_ok=True)
        os.makedirs('./empy/emptytxt', exist_ok=True)
        shutil.copy2(image_path, './empy/empyimg/'+image_path.split('/')[-1])
        shutil.copy2(label_path, './empy/emptytxt/'+label_path.split('/')[-1] )
    if opt.remove_empty:
        os.remove(image_path)
        os.remove(label_path)
def xywhn2xyxy(x, w=640, h=640, padw=0, padh=0):
    y = np.copy(x)
    y[:, 0] = w * (x[:, 0] - x[:, 2] / 2) + padw  # top left x
    y[:, 1] = h * (x[:, 1] - x[:, 3] / 2) + padh  # top left y
    y[:, 2] = w * (x[:, 0] + x[:, 2] / 2) + padw  # bottom right x
    y[:, 3] = h * (x[:, 1] + x[:, 3] / 2) + padh  # bottom right y
    return y
def plot_one_label_on_images(image_path,label_path,opt):
    global EMPTY_TXT_NUMBER
    img = cv2.imread(str(image_path))
    h, w = img.shape[:2]
    with open(label_path, 'r') as f:
        lb = np.array([x.split() for x in f.read().strip().splitlines()], dtype=np.float32)  # labels

        if lb.shape[0]==0:
            EMPTY_TXT_NUMBER+=1

            if opt.show_empty_img:
                cv2.imshow("empty",img)
                cv2.waitKey(0)
                cv2.destroyWindow()
            if opt.remove_empty or opt.copy_empty:
                remove_empty_images_label(image_path, label_path,opt)

    lb[:, 1:] = xywhn2xyxy(lb[:, 1:], w, h, 0, 0)
    for _, x in enumerate(lb):
        class_label = int(x[0])  # class
        cv2.rectangle(img, (int(x[1]), int(x[2])), (int(x[3]), int(x[4])), (0, 255, 255), thickness=4)
        cv2.putText(img, str(class_label), (int(x[1]), int(x[2] - 2)), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=1,
                    color=(0, 255, 255), thickness=4)
    print("empt label number is:",EMPTY_TXT_NUMBER)
    return img


def plot_more_images(base_image_path,base_label_path,class_name,save_res,opt):
    os.makedirs(save_res,exist_ok=True)
    img_list = os.listdir(base_image_path)
    print(img_list)
    for images_file_name in img_list:
        try:
            one_img_path=os.path.join(base_image_path,images_file_name)
            one_txt_path = os.path.join(base_label_path,images_file_name.replace('.jpg','.txt'))
            plot_img = plot_one_label_on_images(one_img_path,one_txt_path,opt)
            cv2.imwrite(one_img_path.replace(base_image_path[:-1],'./result'),plot_img)
        except:
            print("error  !!!!")

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--image_path', nargs='+', type=str, default='yolodata_with_dog/images/val/', help='images path')
    parser.add_argument('--label_path', type=str, default='yolodata_with_dog/labels/val/',help='labels path')
    parser.add_argument('--save_res', type=str, default='./result', help='plot images save path')
    parser.add_argument('--class_name', type=str, default=['classname1','classname2','classnamen'], help='class name list')
    parser.add_argument('--show_empty_img', type=bool, default=False,help='show img if label txt is empty')
    parser.add_argument('--remove_empty', type=bool, default=False, help='remove empty txt file and image')
    parser.add_argument('--copy_empty',type=bool,default=False,help='save empy label and images')
    opt = parser.parse_args()
    print(opt)
    plot_more_images(opt.image_path, opt.label_path, opt.class_name, opt.save_res,opt)