使用Albumentations增强目标检测任务

1、导入所需的库

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

2、定义函数以可视化图像上的边界框和类标签

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)

3、从磁盘加载图像

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

4、用坐标和类标签定义两个边界框

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

# We will use the mapping from category_id to the class name
# to visualize the class label for the bounding box on the image
category_id_to_name = {17: 'cat', 18: 'dog'}

5、使用边界框可视化原始图像

visualize(image, bboxes, category_ids, category_id_to_name)

6、定义一个增强的管道

transform = A.Compose(
    [A.HorizontalFlip(p=0.5)],
    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,
)

7、另一个例子

transform = A.Compose(
    [A.ShiftScaleRotate(p=0.5)],
    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,
)

8、定义一个比较复杂的增强管道

transform = A.Compose([
        A.HorizontalFlip(p=0.5),
        A.ShiftScaleRotate(p=0.5),
        A.RandomBrightnessContrast(p=0.3),
        A.RGBShift(r_shift_limit=30, g_shift_limit=30, b_shift_limit=30, p=0.3),
    ],
    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,
)

min_area 和 min_visibility 参数

• 如果应用空间增强（例如，裁剪图像的一部分或调整图像大小），边界框的大小可能会更改。
• min_area 和 min_visibility 参数控制 Albumentations 在扩充后大小发生变化时应对增强边界框执行的操作。如果应用空间增强（例如，裁剪图像的一部分或调整图像大小），边界框的大小可能会更改。
• min_area 是以像素为单位的值。 如果增强后边界框的面积小于 min_area，Albumentations 将丢弃该框。 因此，返回的增强边界框列表将不包含该边界框。
• min_visibility 是一个介于 0 和 1 之间的值。如果增强后的边界框面积与增强前的边界框面积之比小于 min_visibility，Albumentations 将丢弃该框。 因此，如果扩充过程削减了大部分边界框，则该框将不会出现在扩充边界框的返回列表中。

如果您没有通过 min_area 和 min_visibility 参数，Albumentations 将使用 0 作为它们的默认值。

transform = A.Compose(
    [A.CenterCrop(height=280, width=280, p=1)],
    bbox_params=A.BboxParams(format='coco', label_fields=['category_ids']),
)

transformed = transform(image=image, bboxes=bboxes, category_ids=category_ids)
visualize(
    transformed['image'],
    transformed['bboxes'],
    transformed['category_ids'],
    category_id_to_name,
)

接下来，我们将设置 min_area 值为 4500 像素。

transform = A.Compose(
    [A.CenterCrop(height=280, width=280, p=1)],
    bbox_params=A.BboxParams(format='coco', min_area=4500, label_fields=['category_ids']),
)

transformed = transform(image=image, bboxes=bboxes, category_ids=category_ids)
visualize(
    transformed['image'],
    transformed['bboxes'],
    transformed['category_ids'],
    category_id_to_name,
)

最后，我们将设置 min_visibility 到 0.3。因此，如果输出边界框的面积小于原始区域的 30%，则 Albumentations 将不会返回该边界框。

transform = A.Compose(
    [A.CenterCrop(height=280, width=280, p=1)],
    bbox_params=A.BboxParams(format='coco', min_visibility=0.3, label_fields=['category_ids']),
)

transformed = transform(image=image, bboxes=bboxes, category_ids=category_ids)
visualize(
    transformed['image'],
    transformed['bboxes'],
    transformed['category_ids'],
    category_id_to_name,
)

输出不包含任何边界框。