数据增强：拼接四张小图成为一张大图，同时合并标注信息

目录

0. 背景：

1. 数据准备

2. 合并思路

3. 拼接前后示例

4. 完整代码

0. 背景：

在收集数据时，有时候我们会获取一批尺寸相对较小的图片，其尺寸之小甚至小于我们网络设定的输入尺寸；这时候，如果将多个小图拼接一下，就可以让避免在预处理时无意义的缩放，同时，可以让每个batch中的图片包含更多语义信息，类似于mosaic数据增强。

因此，我决定写个脚本，将这批小图进行每组四张进行组合拼接，并合并每个小图已有的标注信息，避免合并大图重新标注。

1. 数据准备

首先需要使用labelImg工具将小图标注完毕，并将图片、xml文件存放到同一个文件夹下（PS：你可以按照自己的标注格式以及存放习惯，只需要改动代码中的图片、标注文件的读取路径即可）。形如：

2. 合并思路

• 拼接方式选用对角拼接，也即：四张小图各选一个顶点相接，如img1的右下角、img2的左下角、img3的右上角、img4的左上角；
• 随机选定四张小图，用于对角拼接；
• 获取每个小图的标注信息；
• 根据四张小图的尺寸，确定出合并后大图的尺寸及拼接中心点；
• 新建一个空的大图，将各个小图填充进去；
• 保存大图到指定路径；
• 修改每个小图中的标注信息，使之适配大图坐标；
• 合并所有小图的标注信息，形成一个统一的Ground Truth；
• 将统一的GTs写入一个新的xml文件当中，对应所保存的大图；

3. 拼接前后示例

用于拼接的四张小图：

拼接后的大图如下，同时，每个小图的标注信息也会同步到大图当中：

4. 完整代码

"""
Joint small images to be a big image.
"""
import os
import argparse
from PIL import Image
import numpy as np
import xml.etree.ElementTree as ETfrom bboxes2xml import bboxes2xmldef list_dir(path, list_name, suffix='xml'):  # 传入存储的listfor file in os.listdir(path):file_path = os.path.join(path, file)if os.path.isdir(file_path):list_dir(file_path, list_name)else:if file_path.split('.')[-1] == suffix:file_path = file_path.replace('\\', '/')list_name.append(file_path)def get_bboxes(xml_path):tree = ET.parse(open(xml_path, 'rb'))root = tree.getroot()bboxes, cls = [], []for obj in root.iter('object'):obj_cls = obj.find('name').textxmlbox = obj.find('bndbox')xmin = float(xmlbox.find('xmin').text)xmax = float(xmlbox.find('xmax').text)ymin = float(xmlbox.find('ymin').text)ymax = float(xmlbox.find('ymax').text)bboxes.append([xmin, ymin, xmax, ymax])cls.append(obj_cls)bboxes = np.asarray(bboxes, np.int)return bboxes, clsdef main(args):os.makedirs(args.new_data_path, exist_ok=True)imgs = []list_dir(args.raw_data_path, imgs, suffix='jpg')def img_save_name(i, zfill=5):file = args.filename+"_{}.jpg".format(str(i).zfill(zfill))return os.path.join(args.new_data_path, file)for i in range(len(imgs)):# 随机选四张图片, 用于对角拼接img_joint = np.random.choice(imgs, 4, replace=False)img1 = Image.open(img_joint[0])img2 = Image.open(img_joint[1])img3 = Image.open(img_joint[2])img4 = Image.open(img_joint[3])bboxes1, cls1 = get_bboxes(img_joint[0].replace('.jpg', '.xml'))bboxes2, cls2 = get_bboxes(img_joint[1].replace('.jpg', '.xml'))bboxes3, cls3 = get_bboxes(img_joint[2].replace('.jpg', '.xml'))bboxes4, cls4 = get_bboxes(img_joint[3].replace('.jpg', '.xml'))# 大图宽高，及拼接中心点W = max(img1.size[0], img3.size[0]) + max(img2.size[0], img4.size[0])H = max(img1.size[1], img2.size[1]) + max(img3.size[1], img4.size[1])P = [0, 0]P[0] = max(img1.size[0], img3.size[0])P[1] = max(img1.size[1], img2.size[1])# 新建大图，并使用小图填充img_big = Image.new('RGB', (W, H))img_big.paste(img1, (P[0]-img1.size[0], P[1]-img1.size[1]))img_big.paste(img2, (P[0], P[1]-img2.size[1]))img_big.paste(img3, (P[0]-img3.size[0], P[1]))img_big.paste(img4, (P[0], P[1]))img_big.save(img_save_name(i))# 修改每个小图的Bbox，并合并所有子图的bbox和clsbbox_list = []if len(bboxes1) !=0:bboxes1[:, 0], bboxes1[:, 2] = bboxes1[:, 0]+P[0]-img1.size[0], bboxes1[:, 2]+P[0]-img1.size[0]bboxes1[:, 1], bboxes1[:, 3] = bboxes1[:, 1]+P[1]-img1.size[1], bboxes1[:, 3]+P[1]-img1.size[1]bbox_list.append(bboxes1)if len(bboxes2) != 0:bboxes2[:, 0], bboxes2[:, 2] = bboxes2[:, 0]+P[0], bboxes2[:, 2]+P[0]bboxes2[:, 1], bboxes2[:, 3] = bboxes2[:, 1]+P[1]-img2.size[1], bboxes2[:, 3]+P[1]-img2.size[1]bbox_list.append(bboxes2)if len(bboxes3) != 0:bboxes3[:, 0], bboxes3[:, 2] = bboxes3[:, 0]+P[0]-img3.size[0], bboxes3[:, 2]+P[0]-img3.size[0]bboxes3[:, 1], bboxes3[:, 3] = bboxes3[:, 1]+P[1], bboxes3[:, 3]+P[1]bbox_list.append(bboxes3)if len(bboxes4) != 0:bboxes4[:, 0], bboxes4[:, 2] = bboxes4[:, 0]+P[0], bboxes4[:, 2]+P[0]bboxes4[:, 1], bboxes4[:, 3] = bboxes4[:, 1]+P[1], bboxes4[:, 3]+P[1]bbox_list.append(bboxes4)if len(bbox_list) != 0:bboxes = np.vstack(tuple(bbox_list))cls = cls1 + cls2 + cls3 + cls4gts = [[c] + b.tolist() for c, b in zip(cls, bboxes)]bboxes2xml(folder=args.new_data_path.split('/')[-1],img_name=img_save_name(i).replace('\\', '/').split('/')[-1].replace('.jpg', ''),width=W, height=H,gts=gts, xml_save_to=args.new_data_path)def parse_args():parser = argparse.ArgumentParser()parser.add_argument("--raw_data_path", default="raw_data_path", type=str,help="raw dataset files")parser.add_argument("--new_data_path", default="new_data_path", type=str,help="generated new dataset files")parser.add_argument("--filename", default="your_filename", type=str, help="save name")return parser.parse_args()if __name__ == '__main__':args = parse_args()main(args)

其中，bboxes2xml方法定义在bboxes2xml.py文件中：

# bboxes2xml.py filedef bboxes2xml(folder, img_name, width, height, gts, xml_save_to):""" Generate VOC type annotations from prediction resultsArgs:img_name: raw image name without suffix;width, height: size of img;xml_save_to: path of xml will save to;gts: Ground True, [[label, xmin, ymin, xmax, ymax],...,[...]];"""# write in xml filexml_file = open((xml_save_to + '/' + img_name + '.xml'), 'w')xml_file.write('<annotation>\n')xml_file.write('    <folder>' + folder + '</folder>\n')xml_file.write('    <filename>' + str(img_name) + '.jpg' + '</filename>\n')xml_file.write('    <size>\n')xml_file.write('        <width>' + str(width) + '</width>\n')xml_file.write('        <height>' + str(height) + '</height>\n')xml_file.write('        <depth>3</depth>\n')xml_file.write('    </size>\n')# write the region of image on xml filefor gt in gts:xml_file.write('    <object>\n')xml_file.write('        <name>' + str(gt[0]) + '</name>\n')xml_file.write('        <pose>Unspecified</pose>\n')xml_file.write('        <truncated>0</truncated>\n')xml_file.write('        <difficult>0</difficult>\n')xml_file.write('        <bndbox>\n')xml_file.write('            <xmin>' + str(gt[1]) + '</xmin>\n')xml_file.write('            <ymin>' + str(gt[2]) + '</ymin>\n')xml_file.write('            <xmax>' + str(gt[3]) + '</xmax>\n')xml_file.write('            <ymax>' + str(gt[4]) + '</ymax>\n')xml_file.write('        </bndbox>\n')xml_file.write('    </object>\n')xml_file.write('</annotation>')xml_file.close()