前言

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

最近项目遇到的问题是，一批大图已经做好标注，需要将其裁切，同时标签也要进行同步裁切。本文讲解如何实现这一需求，同时将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()