Albumentations概述：用于圖像數(shù)據(jù)增強(qiáng)的開源庫

點(diǎn)擊下方卡片，關(guān)注“新機(jī)器視覺”公眾號(hào)

視覺/圖像重磅干貨，第一時(shí)間送達(dá)

為什么是Albumentations？

• 它是開源的，
• 簡單
• 快速
• 擁有60多種不同的增強(qiáng)
• 有案例
• 而且，最重要的是，可以同時(shí)增強(qiáng)圖像及其掩碼，邊界框或關(guān)鍵點(diǎn)位置。

簡短教程

1. 安裝Albumentations。我真的建議你檢查一下是否有最新的版本，因?yàn)榕f的版本可能有問題。我使用的是“1.0.0”版本，它運(yùn)行良好。
2. 下載下面帶有標(biāo)簽的測試圖像。這只是來自COCO數(shù)據(jù)集的隨機(jī)圖像。我對(duì)它做了一些修改，并以Albumentations要求的格式存儲(chǔ)了它。

import pickle 
import numpy as np 
import matplotlib.pyplot as plt 
import matplotlib.patches as patches

# 加載數(shù)據(jù)
with open("image_data.pickle", "rb") as handle:
    image_data = pickle.load(handle)
    
image = image_data["image"]
mask = image_data["mask"]
bbox = image_data["bbox_coco"]

# 可視化數(shù)據(jù)
fig, ax = plt.subplots(1, 2, figsize=(12, 5))
ax[0].imshow(image)
ax[0].set_title("Image")
ax[1].imshow(image)
bbox_rect = patches.Rectangle(
    bbox[:2], bbox[2], bbox[3], linewidth=2, edgecolor="r", facecolor="none"
)
ax[1].add_patch(bbox_rect)
ax[1].imshow(mask, alpha=0.3, cmap="gray_r")
ax[1].set_title("Image + BBox + Mask")
plt.show()

• 通過使用Compose對(duì)象組合多個(gè)增強(qiáng)來定義轉(zhuǎn)換。
• 每個(gè)增強(qiáng)都有參數(shù)“p”，即要應(yīng)用的概率，另外還有增廣特定的參數(shù)，如RandomCrop的“width”和“height”。
• 使用定義的變換作為函數(shù)來增強(qiáng)圖像及其掩碼。此函數(shù)返回一個(gè)帶有鍵--“image”和“mask”的字典。

import albumentations as A

# 定義增強(qiáng)
transform = A.Compose([
    A.RandomCrop(width=256, height=256, p=1),
    A.HorizontalFlip(p=0.5),
])

# 增強(qiáng)和可視化圖像
fig, ax = plt.subplots(2, 3, figsize=(15, 10))
for i in range(6):
    transformed = transform(image=image, mask=mask)
    ax[i // 3, i % 3].imshow(transformed["image"])
    ax[i // 3, i % 3].imshow(transformed["mask"], alpha=0.3, cmap="gray_r")
plt.show()

• 此外，定義“bbox_params”，其中指定邊界框的格式和邊界框類的參數(shù)coco是指coco數(shù)據(jù)集格式的邊界框-[x_min，y_min，width，height]。參數(shù)'bbox_classes'稍后將用于傳遞邊界框的類。
• transform接受邊界框作為列表列表。此外，即使圖像中只有一個(gè)邊界框，也需要邊界框類（作為列表）。

# 定義增強(qiáng)
transform = A.Compose([
     A.RandomCrop(width=256, height=256, p=1),
     A.HorizontalFlip(p=0.5), 
], bbox_params=A.BboxParams(format='coco', label_fields=["bbox_classes"]))

# 擴(kuò)充和可視化
bboxes = [bbox]
bbox_classes = ["horse"]

fig, ax = plt.subplots(2, 3, figsize=(15, 10))
for i in range(6):
    transformed = transform(
        image=image, 
        bboxes=bboxes, 
        bbox_classes=bbox_classes
    )
    ax[i // 3, i % 3].imshow(transformed["image"])
    trans_bbox = transformed["bboxes"][0]
    bbox_rect = patches.Rectangle(
        trans_bbox[:2],
        trans_bbox[2],
        trans_bbox[3],
        linewidth=2,
        edgecolor="r",
        facecolor="none",
    )
    ax[i // 3, i % 3].add_patch(bbox_rect)
plt.show()

# 定義增強(qiáng)
transform = A.Compose([
     A.RandomCrop(width=256, height=256, p=1),
     A.HorizontalFlip(p=0.5), 
], bbox_params=A.BboxParams(format='coco', label_fields=["bbox_classes"]))

# 增強(qiáng)和可視化
bboxes = [bbox]
bbox_classes = ["horse"]

fig, ax = plt.subplots(2, 3, figsize=(15, 10))
for i in range(6):
    transformed = transform(
        image=image, 
        mask=mask, 
        bboxes=bboxes, 
        bbox_classes=bbox_classes
    )
    ax[i // 3, i % 3].imshow(transformed["image"])
    trans_bbox = transformed["bboxes"][0]
    bbox_rect = patches.Rectangle(
        trans_bbox[:2],
        trans_bbox[2],
        trans_bbox[3],
        linewidth=2,
        edgecolor="r",
        facecolor="none",
    )
    ax[i // 3, i % 3].add_patch(bbox_rect)
    ax[i // 3, i % 3].imshow(transformed["mask"], alpha=0.3, cmap="gray_r")
plt.show()

from torch.utils.data import Dataset
from albumentations.pytorch import ToTensorV2

class CustomDataset(Dataset):
    def __init__(self, images, masks):
        self.images = images  # 假設(shè)這是一個(gè)numpy圖像列表
        self.masks = masks  # 假設(shè)這是一個(gè)numpy掩碼列表
        self.transform = A.Compose([
            A.RandomCrop(width=256, height=256, p=1),
            A.HorizontalFlip(p=0.5),
            ToTensorV2,
        ])
        
    def __len__(self):
        return len(self.images)
        
    def __getitem__(self, idx):
        """返回單個(gè)樣本"""
        image = self.images[idx]
        mask = self.masks[idx]
        transformed = self.transform(image=image, mask=mask)
        transformed_image = transformed["image"]
        transformed_mask = transformed["mask"]
        return transformed_image, transformed_mask

from tensorflow import keras

class CustomDataset(keras.utils.Sequence):
    def __init__(self, images, masks):
        self.images = images
        self.masks = masks
        self.batch_size = 1
        self.img_size = (256, 256)
        self.transform = A.Compose([
            A.RandomCrop(width=256, height=256, p=1), 
            A.HorizontalFlip(p=0.5),
        ])
        
    def __len__(self):
        return len(self.images) // self.batch_size
        
    def __getitem__(self, idx):
        """返回樣本batch"""
        i = idx * self.batch_size
        batch_images = self.images[i : i + self.batch_size]
        batch_masks = self.masks[i : i + self.batch_size]
        batch_images_stacked = np.zeros(
            (self.batch_size,) + self.img_size + (3,), dtype="uint8"
        )
        batch_masks_stacked = np.zeros(
            (self.batch_size,) + self.img_size, dtype="float32"
        )
        for i in range(len(batch_images)):
            transformed = self.transform(
                image=batch_images[i], 
                mask=batch_masks[i]
            )
            batch_images_stacked[i] = transformed["image"]
            batch_masks_stacked[i] = transformed["mask"]
        return batch_images_stacked, batch_masks_stacked