Python深度学习：使用Albumentations进行带边界框的图像增强

环境准备与库导入

为了进行带有边界框的图像增强，首先需要导入必要的Python库。这包括用于图像读取和处理的OpenCV、用于可视化的Matplotlib，以及核心的Albumentations库。

import random
import cv2
from matplotlib import pyplot as plt
import albumentations as A

定义可视化辅助函数

为了直观地检查增强效果，我们需要编写可视化函数。这些函数参考了Detectron的实现，用于在图像上绘制边界框并显示对应的类别标签。

首先定义颜色常量和绘制单个边界框的函数。该函数接收图像、边界框坐标（x_min, y_min, width, height）、类别名称以及颜色和线宽参数。它会在图像上绘制矩形框，并在框的顶部添加带有类别名称的标签。

BOX_COLOR = (255, 0, 0) # Red
TEXT_COLOR = (255, 255, 255) # White

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)

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

    ((text_width, text_height), _) = cv2.getTextSize(class_name, cv2.FONT_HERSHEY_SIMPLEX, 0.35, 1)    
    cv2.rectangle(img, (x_min, y_min - int(1.3 * text_height)), (x_min + text_width, y_min), BOX_COLOR, -1)
    cv2.putText(
        img,
        text=class_name,
        org=(x_min, y_min - int(0.3 * text_height)),
        fontFace=cv2.FONT_HERSHEY_SIMPLEX,
        fontScale=0.35, 
        color=TEXT_COLOR, 
        lineType=cv2.LINE_AA,
    )
    return img

def visualize(image, bboxes, category_ids, category_id_to_name):
    img = image.copy()
    for bbox, category_id in zip(bboxes, category_ids):
        class_name = category_id_to_name[category_id]
        img = visualize_bbox(img, bbox, class_name)
    plt.figure(figsize=(12, 12))
    plt.axis('off')
    plt.imshow(img)

加载图像与标注数据

接下来，从示例数据集中加载图像。这里以COCO数据集为例，使用特定的图像ID（如386298）读取图像文件，并将其从BGR格式转换为RGB格式，以便与Matplotlib兼容。

image = cv2.imread('images/000000386298.jpg')
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

同时，定义边界框坐标、类别ID以及类别ID到名称的映射。COCO格式的边界框由[x_min, y_min, width, height]四个值描述。例如，这里定义了两个边界框和对应的类别ID（17代表猫，18代表狗）。

bboxes = [[5.66, 138.95, 147.09, 164.88], [366.7, 80.84, 132.8, 181.84]]
category_ids = [17, 18]

category_id_to_name = {17: 'cat', 18: 'dog'}

调用可视化函数，查看原始图像及其边界框标注情况。

visualize(image, bboxes, category_ids, category_id_to_name)

构建增强管道

图像增强的核心在于定义变换管道。在目标检测任务中，直接进行随机裁剪可能会导致部分边界框被裁剪掉，从而丢失标注信息。Albumentations提供了RandomSizedBBoxSafeCrop变换来解决这一问题。

该变换会随机裁剪图像的一部分，但会确保裁剪后的区域包含原始图像中的所有边界框。如果无法完全包含所有边界框，它会根据erosion_rate参数允许一定比例的边界框面积丢失。例如，设置erosion_rate=0.2表示增强后的图像中，最多允许20%的边界框面积丢失。

在定义管道时，需要指定输出图像的宽度和高度（如448x336），并配置BboxParams以指定边界框的格式（如COCO格式）和标签字段（如category_ids）。

transform = A.Compose(
    [A.RandomSizedBBoxSafeCrop(width=448, height=336, erosion_rate=0.2)],
    bbox_params=A.BboxParams(format='coco', label_fields=['category_ids']),
)

应用增强与结果验证

为了演示增强效果，可以固定随机种子以确保每次运行产生相同的结果。需要注意的是，在实际的训练管道中，不应在应用转换前固定随机种子，否则会导致每次增强都产生相同的图像，失去增强的意义。

random.seed(7)
transformed = transform(image=image, bboxes=bboxes, category_ids=category_ids)
visualize(
    transformed['image'],
    transformed['bboxes'],
    transformed['category_ids'],
    category_id_to_name,
)

通过改变随机种子，可以得到不同的增强结果。例如，设置种子为3、444等，可以观察到不同的裁剪和缩放效果，但所有边界框都会被尽量保留在图像中。

random.seed(3)
transformed = transform(image=image, bboxes=bboxes, category_ids=category_ids)
visualize(
    transformed['image'],
    transformed['bboxes'],
    transformed['category_ids'],
    category_id_to_name,
)

random.seed(444)
transformed = transform(image=image, bboxes=bboxes, category_ids=category_ids)
visualize(
    transformed['image'],
    transformed['bboxes'],
    transformed['category_ids'],
    category_id_to_name,
)

总结

本文详细讲解了如何在Python深度学习项目中利用Albumentations库进行图像增强。通过定义可视化函数、加载COCO格式标注数据，并应用RandomSizedBBoxSafeCrop变换，我们成功实现了在随机裁剪和缩放图像时保留所有边界框的目标。这种方法对于目标检测等需要精确标注信息的任务尤为重要。

以上就是Python深度学习使用Albumentations进行带边界框的图像增强的详细内容，更多关于Python,深度学习,计算机视觉的资料请关注本站其它相关文章！