【目标检测】大图包括标签切分，并转换成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 tqdm
rootdir = r"E:\Dataset\数据集\可见光数据\原始未裁剪\img"
savedir = r'E:\Dataset\数据集\可见光数据\裁剪后数据\img' # 保存图片文件夹
dis = 1280
leap = 1280
def main():
# 创建输出文件夹
if Path(savedir).exists():
shutil.rmtree(savedir)
os.mkdir(savedir)
num_dir = len(os.listdir(rootdir)) # 得到文件夹下数量
num = 0
for 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 = 0
for 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) / 2
ycenter = ymin + (ymax - ymin) / 2
w = xmax - xmin
h = 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 etree
xml_file_path = "E:/Dataset/数据集/可见光数据/原始未裁剪/labels"
output_txt_path = "E:/Dataset/数据集/可见光数据/裁剪后数据/labels"
class_dict = ['class1', 'class2']
leap = 1280
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 xml2txt(xmin, xmax, ymin, ymax, class_name, img_width, img_height):
# 类别索引
class_index = class_dict.index(class_name)
# 将box信息转换到yolo格式
xcenter = xmin + (xmax - xmin) / 2
ycenter = ymin + (ymax - ymin) / 2
w = xmax - xmin
h = 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
def 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 = leap
img_width = leap
img_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 % 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 = 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()本文参与 腾讯云自媒体分享计划，分享自作者个人站点/博客。
原始发表：2023-10-08，如有侵权请联系 cloudcommunity@tencent.com 删除
评论
登录后参与评论
推荐阅读
目录