目标检测-labelme/labelimg注释文件xml处理

        由于labelme中对x和y的定义与opencv/numpy数值计算不同（如下图），因此，在进行bbox计算时需注意xml文件中的x=w且y=h，此外xml文件中的ncols始终为最长的那条边，即

图片竖直情况：x=w=nrows，y=h=ncols；

图片横平情况：x=w=ncols，y=h=nrows；

1、xml生成coco格式

        生成yolo format数据，读者可根据自己模型的需求进行调整。

        yolo format：1）每个图片对应一个.txt文件；

                             2）文件中一个object一行，每行信息包括class x_center y_center width height；

                             3）坐标取值要正则化，处理到[0,1]范围；

                             4）class类别取值从0开始。

# -*- coding: utf-8 -*-
"""
Created on Tue Sep 17 11:12:05 2022

Usage:
    python coco.py --image_dir data/coco/images --label_dir data/coco/annotations --coco_dir data/coco/labels

@author: ****
"""

import os
import cv2
import math
import glob
import argparse
import pandas as pd
import numpy as np
import xml.etree.ElementTree as ET

objects = {
    'label1':0,
    'label2':1
}

def labelme_xml_to_coco(image_dir, xml_dir, coco_dir):
    for image_path in glob.glob(image_dir + '/*.jpg'):
        image_name = os.path.basename(image_path)
        xml_name = image_name.replace('jpg','xml')
        xml_path = os.path.join(xml_dir, xml_name)
        if not os.path.exists(xml_path):
            # if no objects in image, no *.txt file is required
            pass
        else:
            coco_name = image_name.replace('jpg','txt')
            output = open(os.path.join(coco_dir, coco_name),'w+')

            img = cv2.imread(image_path)
            h,w,_ = img.shape
            print(h,w)

            tree = ET.parse(xml_path)
            root = tree.getroot()
            for member in root.findall('object'):
                object_name = member.find('name').text
                x_list = []
                y_list = []
                # labelme中x为opencv中的w，y为opencv中的h
                for x in member.findall('./polygon/pt/x'):
                    y_list.append(int(float(x.text)))
                for y in member.findall('./polygon/pt/y'):
                    x_list.append(int(float(y.text)))
                # Box coordinates must be in normalized xywh format (from 0 - 1)
                x_left_top = min(x_list) / h
                y_left_top = min(y_list) / w
                x_right_down = max(x_list) / h
                y_right_down = max(y_list) / w

                center_x = (x_right_down + x_left_top)/2
                center_y = (y_right_down + y_left_top) / 2
                object_w = x_right_down - x_left_top
                object_h = y_right_down - y_left_top
                object_class = objects[object_name]
                output.write('%s\t%s\t%s\t%s\t%s\n' % (object_class, center_x, center_y,object_w,object_h))
            output.close()


parser = argparse.ArgumentParser(description='xml_to_coco, a tool to extract box.')
parser.add_argument('--image_dir', help='dir of image files')
parser.add_argument('--label_dir', help='dir of labelme annotated files')
parser.add_argument('--coco_dir', help='dir of box txt output')
args = parser.parse_args()

if not args.image_dir:
    print('You must supply a image_dir\n')
    parser.print_help()
    exit(0)

if not args.label_dir:
    print('You must supply a label_dir\n')
    parser.print_help()
    exit(0)

if not args.coco_dir:
    print('You must supply a coco_dir\n')
    parser.print_help()
    exit(0)

if __name__ == '__main__':
    image_dir = args.image_dir
    xml_dir = args.label_dir
    coco_dir = args.coco_dir
    labelme_xml_to_coco(image_dir,xml_dir,coco_dir)

2、xml文件生成csv格式

# -*- coding: utf-8 -*-
"""
Created on Tue Sep 10 14:40:25 2019

Usage:
    #for labelme:
    python xml-to-csv.py --dir data/val_dataset/xml data/val_dataset/image --label_type labelme
    #for labelimg:
    python xml-to-csv.py --dir data/val_dataset/xml --label_type labelimg
    
@author: ****
"""

import os
import cv2
import math
import glob
import argparse
import pandas as pd
import numpy as np
import xml.etree.ElementTree as ET
 
# source and credits:
# https://raw.githubusercontent.com/datitran/raccoon_dataset/master/xml_to_csv.py
 
def labelimg_xml_to_csv(path, size):
    xml_list = []
    for xml_file in glob.glob(path + '/*.xml'):
        tree = ET.parse(xml_file)
        root = tree.getroot()
        for member in root.findall('object'):
            value = (root.find('filename').text,
                     int(root.find('size')[0].text),
                     int(root.find('size')[1].text),
                     member[0].text,
                     int(member[4][0].text),
                     int(member[4][1].text),
                     int(member[4][2].text),
                     int(member[4][3].text)
                     )
            xml_list.append(value)
    column_name = ['filename', 'width', 'height', 'class', 'xmin', 'ymin', 'xmax', 'ymax']
    xml_df = pd.DataFrame(xml_list, columns=column_name)
    return xml_df

def labelme_xml_to_csv(path, size):
    xml_path,image_path=path[0],path[1]
    xml_list = []
    for xml_file in glob.glob(xml_path + '/*.xml'):
        tree = ET.parse(xml_file)
        root = tree.getroot()
        imagefile = root.find('filename').text
        if os.path.exists(os.path.join(image_path, imagefile)):
            img = cv2.imread(os.path.join(image_path, imagefile))
            shape = np.shape(img)
            x_times = 1.0 if size == False else eval(size)[0]/float(shape[1])
            y_times = 1.0 if size == False else eval(size)[1]/float(shape[0])
            triangle = {}
            for member in root.findall('object'):
                name = member.find('name').text.split('-')[0]

                x_list = []
                y_list = []
                for x in member.findall('./polygon/pt/x'):
                    x_list.append(int(float(x.text)))
                for y in member.findall('./polygon/pt/y'):
                    y_list.append(int(float(y.text)))

                triangle[name] = np.array([x_list, y_list])
            class_keys = triangle.keys()
            for key in ['liaocao', 'liao', 'zhu']:
                if key in class_keys:
                    value = (imagefile,
                             800,
                             1000,
                             key,
                             math.floor(min(triangle[key][0])*x_times),
                             math.floor(min(triangle[key][1])*y_times),
                             math.ceil(max(triangle[key][0])*x_times),
                             math.ceil(max(triangle[key][1])*y_times)
                             )
                    xml_list.append(value)
    column_name = ['filename', 'width', 'height', 'class', 'xmin', 'ymin', 'xmax', 'ymax']
    xml_df = pd.DataFrame(xml_list, columns=column_name)
    return xml_df



def train(path,xml_to_csv,output_path, size):
    xml_df = xml_to_csv(path, size)
    xml_df.to_csv(output_path, index=None)
    print('> tf_wider - Successfully converted xml to csv.')
 

parser = argparse.ArgumentParser(description='xml_to_csv, a tool to extract box.')
parser.add_argument('--dir',nargs='+',help='List or String. if labelme,dir=["xml_dir","image_dir"];\
                    if labelimg,dir=xml_dir')
parser.add_argument('--label_type',help='labelme or labelimg')
parser.add_argument('--csv_path',help='path of csv output')
parser.add_argument('--size',default=False, help='size of resized image')
args = parser.parse_args()

if not args.dir:
    print('You must supply a dir\n')
    parser.print_help()
    exit(0)

if not args.label_type:
    print('You must supply a label_type\n')
    parser.print_help()
    exit(0)

if not args.csv_path:
    print('You must supply a csv_path\n')
    parser.print_help()
    exit(0)

if not args.size:
    print('You must supply a size\n')
    parser.print_help()
    exit(0)

    
if __name__ == '__main__':
    path=args.dir
    label_type=args.label_type
    output_path=args.csv_path
    xml_to_csv=labelme_xml_to_csv if label_type=='labelme' else labelimg_xml_to_csv
    train(path=path,xml_to_csv=xml_to_csv,output_path=output_path, size=args.size)