LabelImg标注的xml格式转yolov5

Posted 2022-08-05 AI浩

import os
import shutil
import numpy as np
import json
from glob import glob
import cv2
from sklearn.model_selection import train_test_split
from os import getcwd
import xml.etree.ElementTree as ET
def convert(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x = (box[0] + box[1]) / 2.0 - 1
    y = (box[2] + box[3]) / 2.0 - 1
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)


def change_2_yolo5(files, txt_Name):
    classes = ["No"]
    imag_name=[]
    for json_file_ in files:
        xml_filename = labelme_path + json_file_ + ".xml"
        out_file = open('%s/%s.txt' % (labelme_path, json_file_), 'w')

        # image_path = labelme_path + json_file['imagePath']

        height, width, channels = cv2.imread(labelme_path + json_file_ + ".jpg").shape
        in_file=open(xml_filename,encoding='utf-8')
        tree = ET.parse(in_file)
        root = tree.getroot()
        imag_name.append(root.find('filename').text)
        for obj in root.iter('object'):
            xmlbox = obj.find('bndbox')
            cls = obj.find('name').text
            if cls not in classes :
                continue
            cls_id = classes.index(cls)
            xmin = int(float(xmlbox.find('xmin').text)) if int(float(xmlbox.find('xmin').text))> 0 else 0
            xmax =int(float(xmlbox.find('xmax').text)) if int(float(xmlbox.find('xmax').text)) > 0 else 0
            ymin = int(float(xmlbox.find('ymin').text)) if int(float(xmlbox.find('ymin').text)) > 0 else 0
            ymax =int(float(xmlbox.find('ymax').text)) if int(float(xmlbox.find('ymax').text)) > 0 else 0
            b = (float(xmin), float(xmax), float(ymin), float(ymax))
            bb = convert((width, height), b)
            out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\\n')
    return imag_name

def image_txt_copy(files,scr_path,dst_img_path,dst_txt_path):
    """
    :param files: 图片名字组成的list
    :param scr_path: 图片的路径
    :param dst_img_path: 图片复制到的路径
    :param dst_txt_path: 图片对应的txt复制到的路径
    :return:
    """
    for file in files:
        img_path=scr_path+file
        shutil.copy(img_path, dst_img_path+file)
        scr_txt_path=scr_path+file.split('.')[0]+'.txt'
        shutil.copy(scr_txt_path, dst_txt_path + file.split('.')[0]+'.txt')


if __name__ == '__main__':
    classes = ["No"]
    # 1.标签路径
    labelme_path = "Annotations/"
    # 3.获取待处理文件
    files = glob(labelme_path + "*.xml")
    files = [i.replace("\\\\", "/").split("/")[-1].split(".xml")[0] for i in files]
    trainval_files, test_files = train_test_split(files, test_size=0.1, random_state=55)
    # split
    train_files, val_files = train_test_split(trainval_files, test_size=0.1, random_state=55)
    train_name_list=change_2_yolo5(train_files, "train")
    print(train_name_list)
    val_name_list=change_2_yolo5(val_files, "val")
    test_name_list=change_2_yolo5(test_files, "test")
    #创建数据集文件夹。
    file_List = ["train", "val", "test"]
    for file in file_List:
        if not os.path.exists('./VOC/images/%s' % file):
            os.makedirs('./VOC/images/%s' % file)
        if not os.path.exists('./VOC/labels/%s' % file):
            os.makedirs('./VOC/labels/%s' % file)
    image_txt_copy(train_name_list,labelme_path,'./VOC/images/train/','./VOC/labels/train/')
    image_txt_copy(val_name_list, labelme_path, './VOC/images/val/', './VOC/labels/val/')
    image_txt_copy(test_name_list, labelme_path, './VOC/images/test/', './VOC/labels/test/')