因为有一批数据有点儿小，数据质量不佳，为了标注方便使用数据增强将数据固定在1080P，方便标注，

# -*- coding: UTF-8 -*-
"""
@Project ：yolov5_relu_fire_smoke_v1.4 
@IDE     ：PyCharm 
@Author  ：沐枫
@Date    ：2024/4/2 20:28

添加白条，做数据增强，最后所有的图片尺寸固定在1080P
"""
import os
import multiprocessing
from concurrent import futures
from copy import deepcopy

import cv2
import numpy as np
import xml.etree.ElementTree as ET
import xml.dom.minidom as minidom


def xyxy2xywh(x):
    """
    Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] where xy1=top-left, xy2=bottom-right
    :param x:
    :return:
    """
    y = np.copy(x)
    y[..., 0] = (x[..., 0] + x[..., 2]) / 2  # x center
    y[..., 1] = (x[..., 1] + x[..., 3]) / 2  # y center
    y[..., 2] = x[..., 2] - x[..., 0]  # width
    y[..., 3] = x[..., 3] - x[..., 1]  # height
    return y


def xywh2xyxy(x: np.ndarray):
    """
    Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
    :param x:
    :return:
    """
    y = np.copy(x)
    y[..., 0] = x[..., 0] - x[..., 2] / 2  # top left x
    y[..., 1] = x[..., 1] - x[..., 3] / 2  # top left y
    y[..., 2] = x[..., 0] + x[..., 2] / 2  # bottom right x
    y[..., 3] = x[..., 1] + x[..., 3] / 2  # bottom right y
    return y


def decodeVocAnnotation(voc_xml_path, class_index_dict):
    """
    voc数据集格式的文件解析，将一个文件解析成一个list，
    使用空格间隔不同对象
    注意:返回的类别不是整型，而是字符串的类别名称
    注意判断返回值是否为 空，如果是空说明没有目标，是一张背景图
    :param voc_xml_path: xml路径
    :param class_index_dict: 类别字典
    :return: [(cls_index, x1, y1, x2, y2), ...]
    """
    assert voc_xml_path.endswith(".xml"), "voc_xml_path must endswith .xml"

    with open(voc_xml_path, 'r', encoding='utf-8') as xml_file:
        # 打开xml文件，并返回根节点
        root = ET.ElementTree().parse(xml_file)

        # 定义一个列表，专门保存目标
        information = []

        # 查找root节点下所有目标信息
        for obj in root.iter('object'):
            # 目标的名称
            name = obj.find('name').text
            # 目标的bbox坐标，一般voc是保存的corner格式的bbox
            box = obj.find('bndbox')
            xmin = box.find('xmin').text
            ymin = box.find('ymin').text
            xmax = box.find('xmax').text
            ymax = box.find('ymax').text

            # 添加一个目标的信息
            # NOTE:返回值的list
            information.append((class_index_dict[name], int(xmin), int(ymin), int(xmax), int(ymax)))

    return information


def create_voc_xml(image_folder, image_filename, width: int, height: int, labels,
                   save_root, class_name_dict, conf_thresh_dict=None):
    """

    :param image_folder: 图片的相对路径
    :param image_filename: 000001.jpg
    :param width: 图片宽
    :param height: 图片高
    :param labels: 目标框:[[class_index, xmin, ymin, xmax, ymax], ...]
    :param save_root: 保存xml的根目录
    :param class_name_dict: cls_index:cls_name,根据index获取正确的类别name
    :param conf_thresh_dict: cls_index:conf_thresh,根据不同类别设置的阈值获取对应的目标，如果设置为None，则表示保存的txt没有置信度
    :return:
    """
    # 创建 XML 文件的根元素
    root = ET.Element("annotation")

    # 添加图片信息
    folder = ET.SubElement(root, "folder")
    folder.text = str(image_folder)
    # 图片名字
    filename = ET.SubElement(root, "filename")
    filename.text = os.path.join(image_filename)
    # 图片大小
    size = ET.SubElement(root, "size")
    width_element = ET.SubElement(size, "width")
    width_element.text = str(width)
    height_element = ET.SubElement(size, "height")
    height_element.text = str(height)
    depth = ET.SubElement(size, "depth")  # 通道数
    depth.text = "3"

    # 添加目标框信息
    for label in labels:
        # 如果该参数设置为None，表示保存的txt没有None
        if conf_thresh_dict is None:
            # 保证这几项是整数
            class_index, x1, y1, x2, y2 = label.astype(dtype=np.int32)
        else:
            class_index, x1, y1, x2, y2, conf = label
            # 保证这几项是整数
            class_index, x1, y1, x2, y2 = np.array([class_index, x1, y1, x2, y2], dtype=np.int32)

            # 根据置信度过滤是否保存项
            if conf < conf_thresh_dict[class_index]:
                continue

        obj = ET.SubElement(root, "object")

        name = ET.SubElement(obj, "name")
        name.text = class_name_dict[int(class_index)]

        pose = ET.SubElement(obj, "pose")
        pose.text = "Unspecified"

        truncated = ET.SubElement(obj, "truncated")
        truncated.text = "0"

        difficult = ET.SubElement(obj, "difficult")
        difficult.text = "0"

        bndbox = ET.SubElement(obj, "bndbox")
        xmin = ET.SubElement(bndbox, "xmin")
        xmin.text = str(x1)

        ymin = ET.SubElement(bndbox, "ymin")
        ymin.text = str(y1)

        xmax = ET.SubElement(bndbox, "xmax")
        xmax.text = str(x2)

        ymax = ET.SubElement(bndbox, "ymax")
        ymax.text = str(y2)

    # 创建 XML 文件并保存
    xml_str = ET.tostring(root, encoding="utf-8")
    xml_str = minidom.parseString(xml_str)
    # 设置缩进为4个空格，xml可读性提高
    pretty_xml = xml_str.toprettyxml(indent=" " * 4)

    save_path = os.path.join(save_root, f"{os.path.splitext(image_filename)[0]}.xml")
    os.makedirs((os.path.dirname(save_path)), exist_ok=True)
    with open(save_path, "w") as xmlFile:
        xmlFile.write(pretty_xml)


def resize_and_pad(image: np.ndarray, labels: np.ndarray, width=1920, height=1080):
    """

    :param image:
    :param labels: (cls_id, x, y, w, h)
    :param width:
    :param height:
    :return:
    """

    def _resize(image: np.ndarray, labels: np.ndarray, width, height):
        """

        :param image:
        :param labels: (cls_id, x, y, w, h)
        :param width:
        :param height:
        :return:
            image: 最后的图片
            labels: (cls_id, x, y, w, h)
        """
        # 判断图片的尺寸，如果尺寸比目标尺寸大，就等比例缩放，斌使用纯白色填充，如果尺寸比目标尺寸小就直接填充到目标尺寸
        img_h, img_w = image.shape[:2]

        if img_w < width and img_h < height:
            # 直接填充
            # 填充的宽度和高度
            dw = (width - img_w) // 2
            dh = (height - img_h) // 2
            # 创建一个新的蒙版
            new_image = np.ones(shape=(height, width, 3), dtype=np.uint8) * 255
            # 将图片填充到里面
            new_image[dh:dh + img_h, dw:dw + img_w, :] = image[:, :, :]

            # 标签平移，(cls_id, x, y, w, h)
            labels[..., 1] += dw
            labels[..., 2] += dh

        else:
            # 等比例缩放后再填充
            # 计算宽度和高度的缩放比例
            ratio = min((width / img_w), (height / img_h))
            # 计算缩放后的宽度和高度
            new_width = int(img_w * ratio)
            new_height = int(img_h * ratio)

            # 等比例缩放图像
            resized_img = cv2.resize(image, (new_width, new_height))

            # 计算需要填充的宽度和高度
            dw = (width - new_width) // 2
            dh = (height - new_height) // 2

            # 创建一个新的蒙版
            new_image = np.ones(shape=(height, width, 3), dtype=np.uint8) * 255
            # 将图片填充到里面
            new_image[dh:dh + new_height, dw:dw + new_width, :] = resized_img[:, :, :]

            # 标签缩放，平移；(cls_id, x, y, w, h)
            labels[..., 1:] *= ratio  # 坐标和宽高都需要缩放
            # 只有中心点需要平移，不影响宽高
            labels[..., 1] += dw
            labels[..., 2] += dh

        return new_image, labels

    SCALE = 2
    # 原图的宽高
    img_h, img_w = image.shape[:2]

    # NOTE:先在外面扩大一次，写和内部函数，在判断，如果图片比目标尺寸大，就等比例缩放，如果图片比目标尺寸小，就直接填充

    # 比较小，先扩大再等比例缩放；比较大，直接等比例缩放
    if img_w < width and img_h < height:
        new_w = img_w * SCALE
        new_h = img_h * SCALE
        # 图片扩大为原来的2倍
        image = cv2.resize(image, (new_w, new_h))
        # labels也扩大,因为图片扩大2倍，所以目标的中心点和宽高都会扩大同样的倍数
        labels[..., 1:] *= SCALE

    # 缩放和填充
    new_image, new_labels = _resize(image, labels, width=width, height=height)

    return new_image, new_labels


def run(image_path, xml_root,
        image_root, save_image_root, save_xml_root,
        class_index_dict, class_name_dict):
    image_file = os.path.basename(image_path)
    image_name, suffix = os.path.splitext(image_file)
    xml_path = image_path.replace(image_root, xml_root).replace(suffix, ".xml")

    if not os.path.exists(xml_path):
        print(f"\n{image_path} no xml\n")
        return

    try:
        # 读图
        image = cv2.imread(image_path)
        if image is None:
            return

        # (cls_id, x, y, w, h)
        labels = decodeVocAnnotation(xml_path, class_index_dict)
        if len(labels) == 0:
            print(f"\n{image_path} no label\n")
            return
        labels = np.array(labels, dtype=np.float64)
        if labels.ndim < 2:
            labels = labels[None, ...]

        # 坐标框转成xywh
        labels[..., 1:] = xyxy2xywh(labels[..., 1:].copy())
        # resize and pad
        new_image, new_labels = resize_and_pad(image, labels.copy(), width=1920, height=1080)
        new_img_h, new_img_w = new_image.shape[:2]
        # 坐标框转成xyxy
        new_labels[..., 1:] = xywh2xyxy(new_labels[..., 1:].copy())

        # 开始保存
        # save_image_path = image_path.replace(image_root, save_image_root)
        save_image_path = os.path.join(save_image_root,
                                       os.path.basename(os.path.dirname(image_path)),
                                       f"aug_{image_file}")
        save_xml_path = save_image_path.replace(save_image_root, save_xml_root).replace(suffix, ".xml")

        os.makedirs(os.path.dirname(save_image_path), exist_ok=True)
        os.makedirs(os.path.dirname(save_xml_path), exist_ok=True)
        # 保存图片
        cv2.imwrite(save_image_path, new_image)
        # 创建xml
        create_voc_xml(image_folder=save_image_path.replace(save_image_root + os.sep, ""),
                       image_filename=os.path.basename(save_image_path),
                       width=new_img_w,
                       height=new_img_h,
                       labels=np.array(new_labels, dtype=np.int32),
                       save_root=os.path.dirname(save_xml_path),
                       class_name_dict=class_name_dict, )
        print(f"\r{image_path}", end='')

    except Exception as e:
        print(f"{image_path} {run.__name__}:{e}")


def run_process(root_file_list,
                image_root, xml_root, save_image_root, save_xml_root,
                class_index_dict, class_name_dict):
    # 使用线程池控制程序执行
    with futures.ThreadPoolExecutor(max_workers=5) as executor:
        for root, file in root_file_list:
            # 向线程池中提交任务，向线程池中提交任务的时候是一个一个提交的
            image_path = os.path.join(root, file)

            executor.submit(run,
                            *(image_path, xml_root,
                              image_root, save_image_root, save_xml_root,
                              class_index_dict, class_name_dict))


if __name__ == '__main__':
    class_index_dict = {
        "fire": 0,
        "smoke": 1,
    }

    class_name_dict = {
        0: "fire",
        1: "smoke",
    }

    data_root = r"Z:\Datasets\FireSmoke_v4"
    data_root = os.path.abspath(data_root)
    # 数据的原图根目录
    image_root = os.path.join(data_root, "images")
    # xml标注文件根目录
    xml_root = os.path.join(data_root, "annotations")
    # 保存根目录
    save_image_root = os.path.join(image_root, "aug-pad")
    save_xml_root = os.path.join(xml_root, "aug-pad")

    # 过滤点不想用的目录
    exclude_dirs = [
        os.sep + r"background",
        os.sep + r"candle_fire",
        os.sep + r"AUG",
        os.sep + r"smoke",
        os.sep + r"val",
        os.sep + r"aug-merge",
        os.sep + r"candle_fire",
        os.sep + r"cut_aug",
        os.sep + r"miniFire",
        os.sep + r"net",
        os.sep + r"new_data",
        os.sep + r"realScenario",
        os.sep + r"smoke",
        os.sep + r"TSMCandle",
    ]

    max_workers = 10  # 线程/进程 数
    print(f"max_workers:{max_workers}")
    # 一个进程处理多少图片
    max_file_num = 3000
    # 保存root和file的list
    root_file_list = list()
    # 创建进程池，根据自己的设备自行调整，别太多，否则会变慢
    pool = multiprocessing.Pool(processes=max_workers)

    for root, _, files in os.walk(image_root):
        # 只要其中有一个，就跳过
        if any(map(lambda x: x in root, exclude_dirs)):
            continue

        for file in files:
            file_name, suffix = os.path.splitext(file)

            if suffix.lower() not in (".jpg", ".jpeg", ".bmp", ".png"):
                continue

            root_file_list.append((root, file))
            if len(root_file_list) > max_file_num:
                pool.apply_async(run_process,
                                 (deepcopy(root_file_list),
                                  image_root, xml_root, save_image_root, save_xml_root,
                                  class_index_dict, class_name_dict))
                # 清除列表中的存储
                root_file_list.clear()
    else:
        pool.apply_async(run_process,
                         (deepcopy(root_file_list),
                          image_root, xml_root, save_image_root, save_xml_root,
                          class_index_dict, class_name_dict))
        # 清除列表中的存储
        root_file_list.clear()

    # 关闭进程池
    pool.close()
    # 等待所有子进程执行结束
    pool.join()

    print("\nFinish ...")