源碼解析一 數據加載 數據增強 datasets.py augmentations.py

參考：https://blog.csdn.net/qq_38253797/article/details/119904518

本文不是簡單的照搬，只是做一個引子，原文中比較詳細，但是很多東西不是非常重要，本文基於yolov6，因此作者將 數據增強部分合並在了augmentations.py中，詳細代碼可在這裏看

1.數據加載

這個函數會在train.py中被調用，用於生成Trainloader, dataset，testloader：

def create_dataloader(path, imgsz, batch_size, stride, single_cls=False,
                      hyp=None, augment=False, cache=False, pad=0.0, rect=False,
                      rank=-1, workers=8, image_weights=False, quad=False, prefix=''):
    """在train.py中被調用，用於生成Trainloader, dataset，testloader
    自定義dataloader函數: 調用LoadImagesAndLabels獲取數據集(包括數據增強) + 調用分佈式採樣器DistributedSampler +
                        自定義InfiniteDataLoader 進行永久持續的採樣數據
    :param path: 圖片數據加載路徑 train/test  如: ../datasets/VOC/images/train2007
    :param imgsz: train/test圖片尺寸（數據增強後大小） 640
    :param batch_size: batch size 大小 8/16/32
    :param stride: 模型最大stride=32   [32 16 8]
    :param single_cls: 數據集是否是單類別 默認False
    :param hyp: 超參列表dict 網絡訓練時的一些超參數，包括學習率等，這裏主要用到裏面一些關於數據增強(旋轉、平移等)的係數
    :param augment: 是否要進行數據增強  True
    :param cache: 是否cache_images False
    :param pad: 設置矩形訓練的shape時進行的填充 默認0.0
    :param rect: 是否開啓矩形train/test  默認訓練集關閉 驗證集開啓
    :param rank:  多卡訓練時的進程編號 rank爲進程編號  -1且gpu=1時不進行分佈式  -1且多塊gpu使用DataParallel模式  默認-1
    :param workers: dataloader的numworks 加載數據時的cpu進程數
    :param image_weights: 訓練時是否根據圖片樣本真實框分佈權重來選擇圖片  默認False
    :param quad: dataloader取數據時, 是否使用collate_fn4代替collate_fn  默認False
    :param prefix: 顯示信息   一個標誌，多爲train/val，處理標籤時保存cache文件會用到
    """

LoadImagesAndLabels

class LoadImagesAndLabels(Dataset):
    # for training/testing
    def __init__(self, path, img_size=640, batch_size=16, augment=False, hyp=None, rect=False,
                 image_weights=False, cache_images=False, single_cls=False, stride=32, pad=0.0, prefix=''):
        """
        初始化過程並沒有什麼實質性的操作,更多是一個定義參數的過程（self參數）,以便在__getitem()__中進行數據增強操作,所以這部分代碼只需要抓住self中的各個變量的含義就算差不多了
        self.img_files: {list: N} 存放着整個數據集圖片的相對路徑
        self.label_files: {list: N} 存放着整個數據集圖片的相對路徑
        cache label -> verify_image_label
        self.labels: 如果數據集所有圖片中沒有一個多邊形label  labels存儲的label就都是原始label(都是正常的矩形label)
                     否則將所有圖片正常gt的label存入labels 不正常gt(存在一個多邊形)經過segments2boxes轉換爲正常的矩形label
        self.shapes: 所有圖片的shape
        self.segments: 如果數據集所有圖片中沒有一個多邊形label  self.segments=None
                       否則存儲數據集中所有存在多邊形gt的圖片的所有原始label(肯定有多邊形label 也可能有矩形正常label 未知數)
        self.batch: 記載着每張圖片屬於哪個batch
        self.n: 數據集中所有圖片的數量
        self.indices: 記載着所有圖片的index
        self.rect=True時self.batch_shapes記載每個batch的shape(同一個batch的圖片shape相同)

collate_fn 注意：這個函數一般是當調用了batch_size次 getitem 函數後纔會調用一次這個函數，對batch_size張圖片和對應的label進行打包。 強烈建議這裏大家debug試試這裏return的數據是不是我說的這樣定義的。

collate_fn4

這裏是yolo-v5作者實驗性的一個代碼 quad-collate function 當train.py的opt參數quad=True 則調用collate_fn4代替collate_fn。 作用：將4張mosaic圖片[1, 3, 640, 640]合成一張大的mosaic圖片[1, 3, 1280, 1280]。將一個batch的圖片每四張處理, 0.5的概率將四張圖片拼接到一張大圖上訓練, 0.5概率直接將某張圖片上採樣兩倍訓練。

img2label_paths

這個文件是根據數據集中所有圖片的路徑找到數據集中所有labels對應的路徑。用在LoadImagesAndLabels模塊的__init__函數中。

verify_image_label

這個函數用於檢查每一張圖片和每一張label文件是否完好。

圖片文件: 檢查內容、格式、大小、完整性

label文件: 檢查每個gt必須是矩形(每行都得是5個數 class+xywh) + 標籤是否全部>=0 + 標籤座標xywh是否歸一化 + 標籤中是否有重複的座標

其他不常用的函數(在detect.py需要使用)：

LoadWebcam：webcam數據流

LoadStreams：支持rtsp流推送

load_image：圖片或者視頻的文件夾載入

工具函數：

可以拆分出來自己使用，也可以在LoadImagesAndLabels函數中，編排自己的規則

flatten_recursive

這個模塊是將一個文件路徑中的所有文件複製到另一個文件夾中 即將image文件和label文件放到一個新文件夾中

extract_boxes

把目標檢測數據集中的每一個gt拆解開 分類別存儲到對應的文件當中，

如果目標檢測數據集之前以及混合過，那麼可以通過這個函數將其分開，不需要手動編寫函數

autosplit

這個模塊是進行自動劃分數據集，weights=(0.9, 0.1, 0.0)爲控制比例關係，

當然，這些函數的初衷是爲，coco數據集準備的。

dataset_stats

這個模塊是統計數據集的信息返回狀態字典。包含: 每個類別的圖片數量 + 每個類別的實例數量

使用建議：特別是在我們拿到數據集以後，需要了解數據類別的分佈關係，經常需要手動編寫代碼查看分類，這個代碼，可以直接統計總類別數目，

在最新的版本中，數據的統計圖是通過這個函數得了的，缺點是比較小，最後會生成統計的json文件，帶有自動下載工具

使用mosaic

在LoadImagesAndLabels函數中，對載入的圖片算法可以使用load_mosaic load_mosaic9 l即爲更改載入策略

load_mosaic函數是拼接四張圖，而load_mosaic9函數是拼接九張圖

這個函數是根據圖片index，從self或者從對應圖片路徑中載入對應index的圖片 並將原圖中hw中較大者擴展到self.img_size, 較小者同比例擴展。
會被用在LoadImagesAndLabels模塊的__getitem__函數和load_mosaic模塊中載入對應index的圖片

2. 數據增強

datasets.py涉及的數據增強在augmentations 總定義，但導入的這些函數增強方法的配置參數是通過配置文件在這裏導入的

from utils.augmentations import Albumentations, augment_hsv, copy_paste, letterbox, mixup, random_perspective

例如：augment_hsv 參數的導入（參數只是作爲響應的概率，並不一定作用）

圖片的色域增強模塊，圖片並不發生移動，所有不需要改變label，只需要 img 增強即可

if self.augment:
            # Albumentations
            img, labels = self.albumentations(img, labels)
            nl = len(labels)  # update after albumentations
            # HSV color-space
            augment_hsv(img, hgain=hyp['hsv_h'], sgain=hyp['hsv_s'], vgain=hyp['hsv_v'])
            # Flip up-down
            if random.random() < hyp['flipud']:
            # Flip left-right
            if random.random() < hyp['fliplr']:

random_perspective

1.這個函數是進行隨機透視變換，對mosaic整合後的圖片進行隨機旋轉、縮放、平移、裁剪，透視變換，並resize爲輸入大小img_size。

2.在mosaic操作之後進行透視變換/仿射變換：

box_candidates

這個函數用在random_perspective中，是對透視變換後的圖片label進行篩選，去除被裁剪過小的框(面積小於裁剪前的area_thr)

還有長和寬必須大於wh_thr個像素，且長寬比範圍在(1/ar_thr, ar_thr)之間的限制。

def box_candidates(box1, box2, wh_thr=2, ar_thr=20, area_thr=0.1, eps=1e-16):
    """用在random_perspective中 對透視變換後的圖片label進行篩選
    去除被裁剪過小的框(面積小於裁剪前的area_thr) 還有長和寬必須大於wh_thr個像素，且長寬比範圍在(1/ar_thr, ar_thr)之間的限制
    Compute candidate boxes: box1 before augment, box2 after augment, wh_thr (pixels), aspect_ratio_thr, area_ratio
    :params box1: [4, n]
    :params box2: [4, n]
    :params wh_thr: 篩選條件 寬高閾值
    :params ar_thr: 篩選條件 寬高比、高寬比最大值閾值
    :params area_thr: 篩選條件 面積閾值
    :params eps: 1e-16 接近0的數 防止分母爲0
    :return i: 篩選結果 [n] 全是True或False   使用比如: box1[i]即可得到i中所有等於True的矩形框 False的矩形框全部刪除

replicate

這個函數是隨機偏移標籤中心，生成新的標籤與原標籤結合。可以用在load_mosaic裏在mosaic操作之後 random_perspective操作之前， 作者默認是關閉的， 自己可以實驗一下效果。

def replicate(img, labels):
    """可以用在load_mosaic裏在mosaic操作之後 random_perspective操作之前  作者默認是關閉的 自己可以實驗一下效果
    隨機偏移標籤中心，生成新的標籤與原標籤結合  Replicate labels
    :params img: img4 因爲是用在mosaic操作之後 所以size=[2*img_size, 2*img_size]
    :params labels: mosaic整合後圖片的所有正常label標籤labels4(不正常的會通過segments2boxes將多邊形標籤轉化爲正常標籤) [N, cls+xyxy]
    :return img: img4 size=[2*img_size, 2*img_size] 不過圖片中多了一半的較小gt個數
    :params labels: labels4 不過另外增加了一半的較小label [3/2N, cls+xyxy]

letterbox

對輸入的圖像進行縮放：原始數據的處理，可見在YOLOV5根本無需考慮圖片的縮放問題。

1.load_image將圖片從文件中加載出來，並resize到相應的尺寸（最長邊等於我們需要的尺寸，最短邊等比例縮放）；

2.letterbox將之前resize後的圖片再pad到我們所需要的放到dataloader中（collate_fn函數）的尺寸（矩形訓練要求同一個batch中的圖片的尺寸必須保持一致）；

3.將label從相對原圖尺寸（原文件中圖片尺寸）縮放到相對letterbox pad後的圖片尺寸。因爲前兩部分的圖片尺寸發生了變化，同樣的我們的label也需要發生相應的變化。

cutout

cutout數據增強，給圖片隨機添加隨機大小的方塊噪聲 ，目的是提高泛化能力和魯棒性。來自論文： https://arxiv.org/abs/1708.04552。

mixup

這個函數是進行mixup數據增強：按比例融合兩張圖片。論文：https://arxiv.org/pdf/1710.09412.pdf。

hist_equalize

進行直方圖均衡化