前言
遥感图像比较大,通常需要切分成小块再进行训练,之前写过一篇关于大图裁切和拼接的文章【目标检测】图像裁剪/标签可视化/图像拼接处理脚本,不过当时的工作流是先将大图切分成小图,再在小图上进行标注,于是就不考虑标签变换的问题。
最近项目遇到的问题是,一批大图已经做好标注,需要将其裁切,同时标签也要进行同步裁切。本文讲解如何实现这一需求,同时将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()