【目标检测】大图包括标签切分，并转换成txt格式

前言

遥感图像比较大，通常需要切分成小块再进行训练，之前写过一篇关于大图裁切和拼接的文章【目标检测】图像裁剪/标签可视化/图像拼接处理脚本，不过当时的工作流是先将大图切分成小图，再在小图上进行标注，于是就不考虑标签变换的问题。

最近项目遇到的问题是，一批大图已经做好标注，需要将其裁切，同时标签也要进行同步裁切。本文讲解如何实现这一需求，同时将labelimg直出的xml格式标签转换成yolov5等模型需要的txt标签。

图片裁剪

图片裁剪还是沿用了一套之前博文提到的编码规则，即将图片裁成1280x1280的图像块，裁剪后通过文件名来标记图像块在原始图像中的位置。

import configparser
import shutil
import yaml
import os.path
from pathlib import Path
from PIL import Image
from tqdm import tqdmrootdir = r"E:\Dataset\数据集\可见光数据\原始未裁剪\img"
savedir = r'E:\Dataset\数据集\可见光数据\裁剪后数据\img'  # 保存图片文件夹dis = 1280
leap = 1280def main():# 创建输出文件夹if Path(savedir).exists():shutil.rmtree(savedir)os.mkdir(savedir)num_dir = len(os.listdir(rootdir))  # 得到文件夹下数量num = 0for parent, dirnames, filenames in os.walk(rootdir):  # 遍历每一张图片filenames.sort()for filename in tqdm(filenames):currentPath = os.path.join(parent, filename)suffix = currentPath.split('.')[-1]if suffix == 'jpg' or suffix == 'png' or suffix == 'JPG' or suffix == 'PNG':img = Image.open(currentPath)width = img.size[0]height = img.size[1]i = j = 0for i in range(0, width, leap):for j in range(0, height, leap):box = (i, j, i + dis, j + dis)image = img.crop(box)  # 图像裁剪image.save(savedir + '/' + filename.split(suffix)[0][:-1] + "__" + str(i) + "__" + str(j) + ".jpg")if __name__ == '__main__':main()

标签裁剪

标签读取

首先需要通过lxml库对xml格式的数据进行解析，主要提取两个信息，1是目标类别，2是目标bbox坐标。

通过递归形式，将xml转换成字典形式，然后就可以获取到需要的信息。

def parse_xml_to_dict(xml):"""将xml文件解析成字典形式"""if len(xml) == 0:  # 遍历到底层，直接返回tag对应的信息return {xml.tag: xml.text}result = {}for child in xml:child_result = parse_xml_to_dict(child)  # 递归遍历标签信息if child.tag != 'object':result[child.tag] = child_result[child.tag]else:if child.tag not in result:result[child.tag] = []result[child.tag].append(child_result[child.tag])return {xml.tag: result}def main():xml_path = r"label.xml"with open(xml_path, encoding="utf-8") as fid:xml_str = fid.read()xml = etree.fromstring(xml_str)data = parse_xml_to_dict(xml)["annotation"]for obj in data["object"]:# 获取每个object的box信息xmin = float(obj["bndbox"]["xmin"])xmax = float(obj["bndbox"]["xmax"])ymin = float(obj["bndbox"]["ymin"])ymax = float(obj["bndbox"]["ymax"])class_name = obj["name"]

标签位置重置

由于图像裁剪成小的图像块，标签也要转换成图像块对应的bbox。不过，对于裁剪的图像，存在的一个问题是，如果标签被切分成两半，该如何进行处理。

下面是我的处理思路，通过对图像块的位置编码，可以分成四种情况。

第一种情况，标签四个角全在图像块中，此时不用做过多处理。
(下图仅为示意，实际尺寸比例未精确，黑色为bbox，红色为切割线)

第二种情况，标签被左右裁开。此时，将左右两部分都当作一个label分给相应的图像块。

第三种情况，标签被上下裁开。此时，将上下两部分都当作一个label分给相应的图像块。

第四种情况，标签被四块裁开，此时，每一块都过于细小，对于小目标而言，这种情况比较少见，因此舍弃该标签。

对应代码：

xmin_index = int(xmin / leap)
xmax_index = int(xmax / leap)
ymin_index = int(ymin / leap)
ymax_index = int(ymax / leap)xmin = xmin % leap
xmax = xmax % leap
ymin = ymin % leap
ymax = ymax % leap# 第一种情况，两个点在相同的图像块中
if xmin_index == xmax_index and ymin_index == ymax_index:info = xml2txt(xmin, xmax, ymin, ymax, class_name, img_width, img_height)file_name = img_name + "__" + str(xmin_index * leap) + "__" + str(ymin_index * leap) + ".txt"write_txt(info, file_name)
# 第二种情况，目标横跨左右两幅图
elif xmin_index + 1 == xmax_index and ymin_index == ymax_index:# 保存左半目标info = xml2txt(xmin, leap, ymin, ymax, class_name, img_width, img_height)file_name = img_name + "__" + str(xmin_index * leap) + "__" + str(ymax_index * leap) + ".txt"write_txt(info, file_name)# 保存右半目标info = xml2txt(0, xmax, ymin, ymax, class_name, img_width, img_height)file_name = img_name + "__" + str(xmax_index * leap) + "__" + str(ymax_index * leap) + ".txt"write_txt(info, file_name)
# 第三种情况，目标纵跨上下两幅图
elif xmin_index == xmax_index and ymin_index + 1 == ymax_index:# 保存上半目标info = xml2txt(xmin, xmax, ymin, leap, class_name, img_width, img_height)file_name = img_name + "__" + str(xmin_index * leap) + "__" + str(ymin_index * leap) + ".txt"write_txt(info, file_name)# 保存下半目标info = xml2txt(xmin, xmax, 0, ymax, class_name, img_width, img_height)file_name = img_name + "__" + str(xmin_index * leap) + "__" + str(ymax_index * leap) + ".txt"write_txt(info, file_name)

标签转换成txt格式

xml格式是 xmin，ymin，xmax，ymax，对应左上角和左下角矩形框的全局像素点坐标。
txt格式是 class， xcenter, ycenter, w, h， 对应中心点和bbox的宽和高，不过该坐标是相对坐标，这里转换时需要除以小图的宽高。

相关代码：

def xml2txt(xmin, xmax, ymin, ymax, class_name, img_width, img_height):# 类别索引class_index = class_dict.index(class_name)# 将box信息转换到yolo格式xcenter = xmin + (xmax - xmin) / 2ycenter = ymin + (ymax - ymin) / 2w = xmax - xminh = ymax - ymin# 绝对坐标转相对坐标，保存6位小数xcenter = round(xcenter / img_width, 6)ycenter = round(ycenter / img_height, 6)w = round(w / img_width, 6)h = round(h / img_height, 6)info = [str(i) for i in [class_index, xcenter, ycenter, w, h]]return info

完整代码

最后附上批量处理的完整代码：

import os
from tqdm import tqdm
from lxml import etreexml_file_path = "E:/Dataset/数据集/可见光数据/原始未裁剪/labels"
output_txt_path = "E:/Dataset/数据集/可见光数据/裁剪后数据/labels"class_dict = ['class1', 'class2']
leap = 1280def parse_xml_to_dict(xml):"""将xml文件解析成字典形式"""if len(xml) == 0:  # 遍历到底层，直接返回tag对应的信息return {xml.tag: xml.text}result = {}for child in xml:child_result = parse_xml_to_dict(child)  # 递归遍历标签信息if child.tag != 'object':result[child.tag] = child_result[child.tag]else:if child.tag not in result:result[child.tag] = []result[child.tag].append(child_result[child.tag])return {xml.tag: result}def xml2txt(xmin, xmax, ymin, ymax, class_name, img_width, img_height):# 类别索引class_index = class_dict.index(class_name)# 将box信息转换到yolo格式xcenter = xmin + (xmax - xmin) / 2ycenter = ymin + (ymax - ymin) / 2w = xmax - xminh = ymax - ymin# 绝对坐标转相对坐标，保存6位小数xcenter = round(xcenter / img_width, 6)ycenter = round(ycenter / img_height, 6)w = round(w / img_width, 6)h = round(h / img_height, 6)info = [str(i) for i in [class_index, xcenter, ycenter, w, h]]return infodef write_txt(info, file_name):with open(file_name, encoding="utf-8", mode="a") as f:# 若文件不为空，添加换行if os.path.getsize(file_name):f.write("\n" + " ".join(info))else:f.write(" ".join(info))def main():for xml_file in os.listdir(xml_file_path):with open(os.path.join(xml_file_path, xml_file), encoding="utf-8") as fid:xml_str = fid.read()xml = etree.fromstring(xml_str)data = parse_xml_to_dict(xml)["annotation"]# img_height = int(data["size"]["height"])# img_width = int(data["size"]["width"])img_height = leapimg_width = leapimg_name = xml_file[:-4]for obj in data["object"]:# 获取每个object的box信息xmin = float(obj["bndbox"]["xmin"])xmax = float(obj["bndbox"]["xmax"])ymin = float(obj["bndbox"]["ymin"])ymax = float(obj["bndbox"]["ymax"])class_name = obj["name"]xmin_index = int(xmin / leap)xmax_index = int(xmax / leap)ymin_index = int(ymin / leap)ymax_index = int(ymax / leap)xmin = xmin % leapxmax = xmax % leapymin = ymin % leapymax = ymax % leap# 第一种情况，两个点在相同的图像块中if xmin_index == xmax_index and ymin_index == ymax_index:info = xml2txt(xmin, xmax, ymin, ymax, class_name, img_width, img_height)file_name = output_txt_path + "/" + img_name + "__" + str(xmin_index * leap) + "__" + str(ymin_index * leap) + ".txt"write_txt(info, file_name)# 第二种情况，目标横跨左右两幅图elif xmin_index + 1 == xmax_index and ymin_index == ymax_index:# 保存左半目标info = xml2txt(xmin, leap, ymin, ymax, class_name, img_width, img_height)file_name = output_txt_path + "/" + img_name + "__" + str(xmin_index * leap) + "__" + str(ymax_index * leap) + ".txt"write_txt(info, file_name)# 保存右半目标info = xml2txt(0, xmax, ymin, ymax, class_name, img_width, img_height)file_name = output_txt_path + "/" + img_name + "__" + str(xmax_index * leap) + "__" + str(ymax_index * leap) + ".txt"write_txt(info, file_name)# 第三种情况，目标纵跨上下两幅图elif xmin_index == xmax_index and ymin_index + 1 == ymax_index:# 保存上半目标info = xml2txt(xmin, xmax, ymin, leap, class_name, img_width, img_height)file_name = output_txt_path + "/" + img_name + "__" + str(xmin_index * leap) + "__" + str(ymin_index * leap) + ".txt"write_txt(info, file_name)# 保存下半目标info = xml2txt(xmin, xmax, 0, ymax, class_name, img_width, img_height)file_name = output_txt_path + "/" + img_name + "__" + str(xmin_index * leap) + "__" + str(ymax_index * leap) + ".txt"write_txt(info, file_name)if __name__ == "__main__":main()