目录
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 ET
from bboxes2xml import bboxes2xml
def list_dir(path, list_name, suffix='xml'): # 传入存储的list
for 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').text
xmlbox = 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, cls
def 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和cls
bbox_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 + cls4
gts = [[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 file
def bboxes2xml(folder, img_name, width, height, gts, xml_save_to):
""" Generate VOC type annotations from prediction results
Args:
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 file
xml_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 file
for 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()