从.xml标注文件voc格式数据中提取bbox，通过yolov4 mosaic数据增强，四张图片合成一张，并生成相应voc格式

import os
import cv2
import glob
import random
import numpy as np
from PIL import Image
import xml.etree.ElementTree as ET
OUTPUT_SIZE = (600, 600)  # Height, Width
SCALE_RANGE = (0.3, 0.7)
FILTER_TINY_SCALE = 1 / 50  # if height or width lower than this scale, drop it.

ANNO_DIR = 'dataset/Annotations/'   #xml path
IMG_DIR = 'dataset/JPEGImages/'     #image path
outdir='img/'
os.makedirs(outdir+'Annotations')
os.makedirs(outdir+'JPEGImages')


xml_head = '''<annotation>
    <folder>VOC2007</folder>
    <filename>{}</filename>.
    <source>
        <database>The VOC2007 Database</database>
        <annotation>PASCAL VOC2007</annotation>
        <image>flickr</image>
        <flickrid>325991873</flickrid>
    </source>
    <owner>
        <flickrid>null</flickrid>
        <name>null</name>
    </owner>    
    <size>
        <width>{}</width>
        <height>{}</height>
        <depth>3</depth>
    </size>
    <segmented>0</segmented>
    '''
xml_obj = '''
    <object>        
        <name>{}</name>
        <pose>Rear</pose>
        <truncated>0</truncated>
        <difficult>0</difficult>
        <bndbox>
            <xmin>{}</xmin>
            <ymin>{}</ymin>
            <xmax>{}</xmax>
            <ymax>{}</ymax>
        </bndbox>
    </object>
    '''
xml_end = '''
</annotation>'''


def main():
    idx_name=1
    img_paths, annos = get_dataset(ANNO_DIR, IMG_DIR)
    idxs = random.sample(range(len(annos)), len(annos))
    len_img=int(len(annos)/4)

    for i in range(len_img):
        print('i-------',i)
        idx_select=[]  # 4n-4
        idx_select.append(idxs[4*i-4]) 
        idx_select.append(idxs[4*i-3]) 
        idx_select.append(idxs[4*i-2]) 
        idx_select.append(idxs[4*i-1]) 
        new_image, new_annos = update_image_and_anno(img_paths, annos,
                                                 idx_select,
                                                 OUTPUT_SIZE, SCALE_RANGE,
                                                 filter_scale=FILTER_TINY_SCALE)
        savename='jfdata11_'+str(i)+'.jpg'
        cv2.imwrite(outdir+'JPEGImages/'+savename, new_image)
        obj = ''
        head = xml_head.format(str(savename),str(600),str(600))
        for anno in new_annos:
            start_point = (int(anno[1] * OUTPUT_SIZE[1]), int(anno[2] * OUTPUT_SIZE[0]))
            end_point = (int(anno[3] * OUTPUT_SIZE[1]), int(anno[4] * OUTPUT_SIZE[0]))
            cv2.rectangle(new_image, start_point, end_point, (0, 255, 0), 1, cv2.LINE_AA)
            #print('anno--------',start_point,end_point)
            category_name=str((anno[0])
            obj += xml_obj.format(category_name,start_point[0],start_point[1],end_point[0],end_point[1])
        cv2.imwrite(outdir+'JPEGImages/output_'+savename, new_image)
  
        new_image = cv2.cvtColor(new_image, cv2.COLOR_BGR2RGB)
        new_image = Image.fromarray(new_image.astype(np.uint8))
        idx_name+=1
        xml_path=outdir+'Annotations/'+savename.replace('.jpg','.xml')#xml save path
        with open(xml_path,'w') as f_xml:
            f_xml.write(head+obj+xml_end) 


def update_image_and_anno(all_img_list, all_annos, idxs, output_size, scale_range, filter_scale=0.):
    output_img = np.zeros([output_size[0], output_size[1], 3], dtype=np.uint8)
    scale_x = scale_range[0] + random.random() * (scale_range[1] - scale_range[0])
    scale_y = scale_range[0] + random.random() * (scale_range[1] - scale_range[0])
    divid_point_x = int(scale_x * output_size[1])
    divid_point_y = int(scale_y * output_size[0])

    new_anno = []
    for i, idx in enumerate(idxs):
        path = all_img_list[idx]
        img_annos = all_annos[idx]

        img = cv2.imread(path)
        if i == 0:  # top-left
            img = cv2.resize(img, (divid_point_x, divid_point_y))
            output_img[:divid_point_y, :divid_point_x, :] = img
            for bbox in img_annos:
                xmin = bbox[1] * scale_x
                ymin = bbox[2] * scale_y
                xmax = bbox[3] * scale_x
                ymax = bbox[4] * scale_y
                new_anno.append([bbox[0], xmin, ymin, xmax, ymax])

        elif i == 1:  # top-right
            img = cv2.resize(img, (output_size[1] - divid_point_x, divid_point_y))
            output_img[:divid_point_y, divid_point_x:output_size[1], :] = img
            for bbox in img_annos:
                xmin = scale_x + bbox[1] * (1 - scale_x)
                ymin = bbox[2] * scale_y
                xmax = scale_x + bbox[3] * (1 - scale_x)
                ymax = bbox[4] * scale_y
                new_anno.append([bbox[0], xmin, ymin, xmax, ymax])
        elif i == 2:  # bottom-left
            img = cv2.resize(img, (divid_point_x, output_size[0] - divid_point_y))
            output_img[divid_point_y:output_size[0], :divid_point_x, :] = img
            for bbox in img_annos:
                xmin = bbox[1] * scale_x
                ymin = scale_y + bbox[2] * (1 - scale_y)
                xmax = bbox[3] * scale_x
                ymax = scale_y + bbox[4] * (1 - scale_y)
                new_anno.append([bbox[0], xmin, ymin, xmax, ymax])
        else:  # bottom-right
            img = cv2.resize(img, (output_size[1] - divid_point_x, output_size[0] - divid_point_y))
            output_img[divid_point_y:output_size[0], divid_point_x:output_size[1], :] = img
            for bbox in img_annos:
                xmin = scale_x + bbox[1] * (1 - scale_x)
                ymin = scale_y + bbox[2] * (1 - scale_y)
                xmax = scale_x + bbox[3] * (1 - scale_x)
                ymax = scale_y + bbox[4] * (1 - scale_y)
                new_anno.append([bbox[0], xmin, ymin, xmax, ymax])

    if 0 < filter_scale:
        new_anno = [anno for anno in new_anno if
                    filter_scale < (anno[3] - anno[1]) and filter_scale < (anno[4] - anno[2])]

    return output_img, new_anno


def get_dataset(anno_dir, img_dir):
    img_paths = []
    annos = []
    for anno_file in glob.glob(os.path.join(anno_dir, '*.xml')):
        anno_id = anno_file.split('/')[-1].split('.xml')[0] # img_name
        img_path = os.path.join(img_dir, f'{anno_id}.jpg')
        img = cv2.imread(img_path)
        img_height, img_width, _ = img.shape
        del img

        boxes = []
        idx=0

        tree = ET.parse(anno_file)
        root = tree.getroot()
        for obj in get(root, 'object'):
            bndbox = get_and_check(obj, 'bndbox', 1)
            class_id = str(get_and_check(obj, 'name', 1).text)
            xmin = int(float(get_and_check(bndbox, 'xmin', 1).text))
            ymin = int(float(get_and_check(bndbox, 'ymin', 1).text))
            xmax = int(float(get_and_check(bndbox, 'xmax', 1).text))
            ymax = int(float(get_and_check(bndbox, 'ymax', 1).text))

            xmin = max(xmin, 0) / img_width
            ymin = max(ymin, 0) / img_height
            xmax = min(xmax, img_width) / img_width
            ymax = min(ymax, img_height) / img_height
            idx+=1
            boxes.append([class_id, xmin, ymin, xmax, ymax])

            if not boxes:
                continue

        img_paths.append(img_path)
        annos.append(boxes)
    return img_paths, annos


def get(root, name):
    return root.findall(name)
 
 
def get_and_check(root, name, length):
    vars = root.findall(name)
    if len(vars) == 0:
        raise NotImplementedError('Can not find %s in %s.'%(name, root.tag))
    if length > 0 and len(vars) != length:
        raise NotImplementedError('The size of %s is supposed to be %d, but is %d.'%(name, length, len(vars)))
    if length == 1:
        vars = vars[0]
    return vars



if __name__ == '__main__':
    main()