參考：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
進行直方圖均衡化