[Deep Learning] YOLO to VOC VOC to YOLO

YOLO转VOC

# coding:utf-8
import os
import os.path
import xml.etree.ElementTree as ET
from copy import copy

import cv2
import numpy as np
from tqdm import tqdm


def convert_annotation(imgpath: str, allImagesFiles: set, txtdir: str, image_id: str, classes: dict, vocxmlpath: str):
    xmlname = os.path.join(vocxmlpath, '%s.xml' % (image_id))

    yolotxt = open(os.path.join(txtdir, '%s.txt' % (image_id)), 'r').readlines()
    if (os.path.join(imgpath, image_id + ".jpg") in allImagesFiles):
        imageName = os.path.join(imgpath, image_id + ".jpg")
    elif (os.path.join(imgpath, image_id + ".png") in allImagesFiles):
        imageName = os.path.join(imgpath, image_id + ".png")
    else:
        return 1
    img = cv2.imdecode(np.fromfile(imageName, dtype=np.uint8), 1)  # img是矩阵
    height, width, depth = img.shape[0], img.shape[1], img.shape[2]

    root = ET.parse("template.xml")
    tree = root.getroot()

    size = tree.find('size')
    size.find('width').text = width
    size.find('height').text = height
    size.find('depth').text = depth

    objtmpl = copy(tree.find('object'))
    tree.remove(tree.find('object'))

    zidian = dict(zip(classes.values(), classes.keys()))
    for obj in yolotxt:
        if len(obj.strip()) > 4:  # 非空
            objtmpl_copy = copy(objtmpl)
            clsint, x, y, w, h = list(map(float, obj.strip().split(" ")))
            objtmpl_copy.find('name').text = zidian[int(clsint)]
            if 0 <= int((x - w / 2) * width) <= width:
                objtmpl_copy.find('bndbox').find('xmin').text = int((x - w / 2) * width)
            elif int((x - w / 2) * width) < 0:
                objtmpl_copy.find('bndbox').find('xmin').text = 0
            else:
                objtmpl_copy.find('bndbox').find('xmin').text = width

            if 0 <= int((y - h / 2) * height) <= height:
                objtmpl_copy.find('bndbox').find('ymin').text = int((y - h / 2) * height)
            elif int((y - h / 2) * height) < 0:
                objtmpl_copy.find('bndbox').find('ymin').text = 0
            else:
                objtmpl_copy.find('bndbox').find('ymin').text = height

            if 0 <= int((x + w / 2) * width) <= width:
                objtmpl_copy.find('bndbox').find('xmax').text = int((x + w / 2) * width)
            elif int((x + w / 2) * width) < 0:
                objtmpl_copy.find('bndbox').find('xmax').text = 0
            else:
                objtmpl_copy.find('bndbox').find('xmax').text = width

            if 0 <= int((y + h / 2) * height) <= height:
                objtmpl_copy.find('bndbox').find('ymax').text = int((y + h / 2) * height)
            elif int((y + h / 2) * height) < 0:
                objtmpl_copy.find('bndbox').find('ymax').text = 0
            else:
                objtmpl_copy.find('bndbox').find('ymax').text = height

            tree.append(objtmpl_copy)

    root.write(xmlname, encoding="utf-8")
    return 0


def listPathAllfiles(dirname):
    result = []
    for maindir, subdir, file_name_list in os.walk(dirname):
        for filename in file_name_list:
            apath = os.path.join(maindir, filename)
            result.append(apath)
    return result


if __name__ == '__main__':
    classes = {
    'idcard': 0}  # 标签名:标签id
    # 给出VOC数据存储地址
    yolotxtpath = r"E:\detection\03idcard\labels"
    imgpath = r"E:\detection\03idcard\images"
    vocxmlpath = r"E:\detection\03idcard\xml"

    allfiles = listPathAllfiles(yolotxtpath)
    allImagesFiles = set(listPathAllfiles(imgpath))
    print("一共有文件个数:", len(allfiles))
    failNum = 0
    for txtName in tqdm(allfiles):
        if txtName.endswith(".txt") and (not txtName.startswith("classes")):
            # xml存储路径,yololabels存储路径,xml文件名称不带.xml后缀,需要的类及其类id的字典
            if convert_annotation(imgpath, allImagesFiles, yolotxtpath, os.path.basename(txtName)[:-4], classes,
                                  vocxmlpath) == 1:
                failNum += 1
    print("失败了多少个文件的labels:", failNum)

VOC转YOLO

# coding:utf-8
import os
import os.path
import xml.etree.ElementTree as ET


def convert_annotation(xmldir: str, txtdir: str, image_id: str, classes: dict):
    in_file = open(os.path.join(xmldir, '%s.xml' % (image_id)), 'r', encoding='UTF-8')
    out_file = open(os.path.join(txtdir, '%s.txt' % (image_id)), 'w')
    tree = ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    size_width = int(size.find('width').text)
    size_height = int(size.find('height').text)
    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult) == 1:
            continue
        cls_id = classes[cls]
        xmlbox = obj.find('bndbox')
        b = [float(xmlbox.find('xmin').text),
             float(xmlbox.find('xmax').text),
             float(xmlbox.find('ymin').text),
             float(xmlbox.find('ymax').text)]

        if size_width == 0 or size_height == 0 or b[0] == b[1] or b[2] == b[3]:
            print("不合理的图不再给labels ", image_id)
            # if os.path.exists(xmldir + '%s.xml' % (image_id)):
            # os.remove(xmldir + '%s.xml' % (image_id))
            out_file.close()
            os.remove(os.path.join(txtdir, '%s.txt' % (image_id)))
            return 1

        # 标注越界修正
        if b[0] < 0:
            b[0] = 0
        if b[1] > size_width:
            b[1] = size_width
        if b[2] < 0:
            b[2] = 0
        if b[3] > size_height:
            b[3] = size_height
        txt_data = [round(((b[0] + b[1]) / 2.0 - 1) / size_width, 6),
                    round(((b[2] + b[3]) / 2.0 - 1) / size_height, 6),
                    round((b[1] - b[0]) / size_width, 6),
                    round((b[3] - b[2]) / size_height, 6)]
        if txt_data[0] < 0 or txt_data[1] < 0 or txt_data[2] < 0 or txt_data[3] < 0:
            print("不合理的图不再给labels ", image_id)
            out_file.close()
            os.remove(os.path.join(txtdir, '%s.txt' % (image_id)))
            return 1
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in txt_data]) + '\n')

    in_file.close()
    out_file.close()
    return 0


def listPathAllfiles(dirname):
    result = []
    for maindir, subdir, file_name_list in os.walk(dirname):
        for filename in file_name_list:
            apath = os.path.join(maindir, filename)
            result.append(apath)
    return result


if __name__ == '__main__':
    classes = {
    'idcard': 0}  # 标签名:标签id

    # 给出VOC数据存储地址
    VOCrootpath = r"E:\detection\04bankcard\hand"
    xmlpath = os.path.join(VOCrootpath, "xml")
    txtpath = os.path.join(VOCrootpath, "labels")

    if not os.path.exists(xmlpath):
        os.makedirs(xmlpath)
    if not os.path.exists(txtpath):
        os.makedirs(txtpath)

    allfiles = listPathAllfiles(xmlpath)
    print("一共有文件个数:", len(allfiles))
    failNum = 0
    for xmlName in allfiles:
        if xmlName.endswith(".xml"):
            # xml存储路径,yololabels存储路径,xml文件名称不带.xml后缀,需要的类及其类id的字典
            if convert_annotation(xmlpath, txtpath, os.path.basename(xmlName)[:-4], classes) == 1:
                failNum += 1
    print("失败了多少个文件的labels:", failNum)