HCHUNGW的部落格

7.6 albumentations 資料增強庫

本節介紹albumentations庫，albumentations是強大的資料增強庫，原計劃在第三章中進行介紹，後因篇幅過大，放到了本章作為進階技巧。

為什麼要用albumentations？ pytorch的transforms有什麼不足麼？ 當然有不足了， pytorch的transforms在處理圖像分割與目標檢測這一類需要圖像與標籤同時變換的時候不方便，而albumentations提供了圖像分割、目標檢測等複雜任務的資料增強方法。

為什麼要用albumentations？

從github中譯過來：

• 支持多種任務：支援分類、語義分割、實例分割、目標檢測和姿態估計等；
• 提供簡潔的API：針對分割、邊界框回歸、關鍵點任務及多種資料形態（RBG-images, grayscale images, multispectral images）均可採用統一的函數完成資料增強。
• 資料增強方法多：提供超70種變換方法。
• 速度快效率高：相比其它常見資料增強庫，多數方法速度都為最優
• 支持主流框架：albumentations已是pytorch生態系統的一員，可很好適配pytorch，同時支援TensorFlow中使用。
• 專家背書：albumentations的作者大多數來自工業界大牛
• 市場驗證：albumentations已經被廣泛應用於工業界和學術界，以及各種競賽，並且獲得了優異成績。

albumentations有那麼多優點，還不趕快來學習它。

安裝

pip install -U albumentations

Copy

上手demo

demo代碼請查看配套代碼，通過代碼可以看到，使用步驟有：

1. 定義一個Compose，內部包括多個變換方法（同transforms一樣）
2. 將compose放到dataset中，在getitem函數中實現調用
3. getitem中注意各變換方法的輸入，由於albumentations支援多種資料同時處理，因此輸入時需要指定變數。如image、mask、bboxes和keypoints之類。
4. 經過Compose返回的資料是一個字典形式，需要根據key獲取對應資訊，如image、mask、bboxes和keypoints。

可以任意調整以下變換方法，以及參數設置，觀察圖像變化。

train_transform = A.Compose(

    [

        A.Resize(512, 512),

        A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.05, rotate_limit=55, p=0.5),

    ]

)

Copy

圖像-標籤對的資料增強

圖像分割、目標檢測和關鍵點任務中的資料增強，需要同時對圖像和標籤一起變換，這也是albumentations與pytorch的transforms最大的區別。

請運行代碼觀察結果，這裡介紹使用步驟

1. dataset的getitem中將image與mask同時傳入self.transfoms()中；
2. 在DataLoader返回的資料中，採用key-value的形式獲取image與mask/bboes/keypoints；

data_augmented = data_transform(image=image_rgb, mask=mask)

Copy

albumentations 代碼結構

albumentations採用BasicTransform作為基類，再根據是否需要對標籤進行變換，劃分為nlyTransform 和 DualTransform

顧名思義，ImageOnlyTransform就是僅針對原圖進行操作，DualTransform就是同時對原圖和標籤進行操作。

兩種類型的變換常見的方如下：

• ImageOnlyTransform： Blur，RGBShift，RandomBrightnessContrast等
• DualTransform： Resize, Flip， Crop， Rotate，ElasticTransform等

為了分析albumentations代碼結構，特地繪製了一副簡略版的UML類圖，再結合Resize與Blur的函數，分析albumentations代碼運行結構。

在代碼中通過Resize類的調試，可看到實現圖像模糊的功能在apply()函數中，接下來就需要理清代碼結構，尋找BasicTransform是如何調用apply()的。

通過代碼單步調試，發現依次進行以下步驟：

1. BasicTransform的__call__()：在89行，return時調用self.apply_with_params(params, **kwargs)
2. BasicTransform的self.apply_with_params()：在100行，通過target_function = self._get_target_function(key)獲得具體需要變換的函數
3. BasicTransform的self._get_target_function()：在123行，通過self.targets.get()獲得具體的target_function
4. DualTransform的targets：是一個字典，其中定義了key與具體變換函數的映射，其中self.apply即上文提到的Resize類下的self.apply。

到這裡一個具體的變換實現過程才算走完，其中BasicTransform定義了基礎邏輯，例如概率選擇、參數獲取等，DualTransform則是定義mask,masks，bboxes，keypoints的變換函數介面，最終的實現由Resize類來完成。同理，ImageOnlyTransform 一樣。

@property

def targets(self):

    return {

        "image": self.apply,

        "mask": self.apply_to_mask,

        "masks": self.apply_to_masks,

        "bboxes": self.apply_to_bboxes,

        "keypoints": self.apply_to_keypoints,

    }

Copy

小結

本節介紹了albumentations的優點，基本使用以及代碼結構，關於albumentations的70多種方法請查看附錄，或者閱讀文檔查看API，使用方法非常簡單，也可到代碼中查看編寫的批量資料增強程式碼片段，實現了68個資料增強視覺化

附錄

albumentations提供70多種API，這裡將不再一一介紹各API的參數，請查看文檔即可。

整體分為兩種，Pixel-Level和Spatial-Level：

• AdvancedBlur
• Blur
• CLAHE
• ChannelDropout
• ChannelShuffle
• ColorJitter
• Downscale
• Emboss
• Equalize
• FDA
• FancyPCA
• FromFloat
• GaussNoise
• GaussianBlur
• GlassBlur
• HistogramMatching
• HueSaturationValue
• ISONoise
• ImageCompression
• InvertImg
• MedianBlur
• MotionBlur
• MultiplicativeNoise
• Normalize
• PixelDistributionAdaptation
• Posterize
• RGBShift
• RandomBrightnessContrast
• RandomFog
• RandomGamma
• RandomRain
• RandomShadow
• RandomSnow
• RandomSunFlare
• RandomToneCurve
• RingingOvershoot
• Sharpen
• Solarize
• Superpixels
• TemplateTransform
• ToFloat
• ToGray
• ToSepia
• UnsharpMask