【深度学习】带标注的数据增强及albumentations包下载

注：代码修改自另一位博主的文章，很抱歉找不到地址了。
做出以下修改：
1、把控制参数放在头部便于修改
2、原代码的增强前后图片和标注存放于一个文件夹，这里存放在新文件夹。
3、做出一些注释
注意：
1、visualize部分未修改，可能会因为我改了变量名跑不通
2、缩放方法可能造成标注偏移，增强之后需要检查
albumentations包下载
如果一直下载不了（time out），用清华镜像源：https://pypi.tuna.tsinghua.edu.cn/simple

pip install albumentations -i https://pypi.tuna.tsinghua.edu.cn/simple

import cv2
from matplotlib import pyplot as plt
import xml.etree.ElementTree as ET
import albumentations as A
import os
import time

# 控制参数
BOX_COLOR = (255, 0, 0)  # Red
TEXT_COLOR = (255, 255, 255)  # White
# 增强张数 original pictures size：62, then total size is 62*GENERATED_PICS_SIZE
GENERATED_PICS_SIZE = 20  # 增强方法在main Compose中修改
# 上级目录
DIR = "D:\\AI\\data6"
# 存储原图片的文件夹名，默认格式未jpg,如果为png需要自行修改
IMAGES_FILE = "images"
# 存储原xml标注的文件夹名
ANNOTATIONS_FILE = "annotations"
# 检查原本的xml标注，object中第几个为bndbox，从0开始计算
OBJ_NUM = 4  # 第五个




def visualize_bbox(img, bbox, class_name, color=BOX_COLOR, thickness=2):
    """Visualizes a single bounding box on the image"""
    # x_min, y_min, w, h = bbox
    # x_min, x_max, y_min, y_max = int(x_min), int(x_min + w), int(y_min), int(
    #     y_min + h)
    x_min, y_min, x_max, y_max = bbox
    print(x_min, y_min, x_max, y_max)

    cv2.rectangle(img, (int(x_min), int(y_min)), (int(x_max), int(y_max)),
                  color=color, t