Tensorflow2.0 實現 YOLOv3（六）：dataset.py

文章說明

本系列文章旨在對 Github 上 malin9402 提供的代碼進行說明，在這篇文章中，我們會對 YOLOv3 項目中的 dataset.py 文件進行說明。在這個文件中，只有一個 Dataset 類，用於生成數據集，所以本文將對代碼的解釋直接寫在代碼旁邊。

如果只是想運行 Github 上的代碼，可以參考對 YOLOv3 代碼的說明一文。

完整代碼

import os
import cv2
import random
import numpy as np
import tensorflow as tf
import core.utils as utils
from core.config import cfg

class Dataset(object):
    """implement Dataset here"""
    def __init__(self, dataset_type):
        self.annot_path  = cfg.TRAIN.ANNOT_PATH if dataset_type == 'train' else cfg.TEST.ANNOT_PATH  # 訓練（測試）集標籤路徑
        self.input_sizes = cfg.TRAIN.INPUT_SIZE if dataset_type == 'train' else cfg.TEST.INPUT_SIZE  # 訓練（測試）集圖片尺寸
        self.batch_size  = cfg.TRAIN.BATCH_SIZE if dataset_type == 'train' else cfg.TEST.BATCH_SIZE  # 訓練（測試）集批次大小
        self.data_aug    = cfg.TRAIN.DATA_AUG   if dataset_type == 'train' else cfg.TEST.DATA_AUG  # 是否對訓練（測試）集圖片進行數據增強處理

        self.train_input_sizes = cfg.TRAIN.INPUT_SIZE  # 訓練集圖片尺寸
        self.strides = np.array(cfg.YOLO.STRIDES)  # 每個 feature map 中的一個格子代表原始圖像中的幾個格子
        self.classes = utils.read_class_names(cfg.YOLO.CLASSES)  # 類別的索引
        self.num_classes = len(self.classes)  # 類別的個數
        self.anchors = np.array(utils.get_anchors(cfg.YOLO.ANCHORS)) # 先驗框的寬度和高度
        self.anchor_per_scale = cfg.YOLO.ANCHOR_PER_SCALE  # 一個尺度（feature map）上有幾個先驗框
        self.max_bbox_per_scale = 150  # 一個尺度（feature map）上最多有幾個真實框（即最多有幾個檢測目標）

        self.annotations = self.load_annotations(dataset_type)  # 加載訓練（測試）集標籤
        self.num_samples = len(self.annotations)  # 樣本數量
        self.num_batchs = int(np.ceil(self.num_samples / self.batch_size))  # 一共有幾個 batch
        self.batch_count = 0  # 計數


    def load_annotations(self, dataset_type):
        with open(self.annot_path, 'r') as f:
            txt = f.readlines()
            annotations = [line.strip() for line in txt if len(line.strip().split()[1:]) != 0]
        np.random.shuffle(annotations)
        return annotations

    def __iter__(self):
        return self

    def __next__(self):

        with tf.device('/cpu:0'):
            self.train_input_size = random.choice(self.train_input_sizes)  # 如果輸入的圖片尺寸不一，這裏可以隨機選擇；但如果輸入的圖片尺寸一致，這裏無論怎麼隨機結果都是一樣的
            self.train_output_sizes = self.train_input_size // self.strides  # 三個輸出（對應三個 feature map）的尺寸

            # 初始化一個批次的樣本
            batch_image = np.zeros((self.batch_size, self.train_input_size, self.train_input_size, 3), dtype=np.float32)

            # 初始化一個批次的輸出（批次大小爲4，三個尺度上的輸出尺寸分別爲52，26，13，一個尺度上有3個先驗框，檢測物共有80類）
            batch_label_sbbox = np.zeros((self.batch_size, self.train_output_sizes[0], self.train_output_sizes[0],
                                          self.anchor_per_scale, 5 + self.num_classes), dtype=np.float32)  # 初始化第一個尺度上的輸出，shape 爲 [4, 52, 52, 3, 85]
            batch_label_mbbox = np.zeros((self.batch_size, self.train_output_sizes[1], self.train_output_sizes[1],
                                          self.anchor_per_scale, 5 + self.num_classes), dtype=np.float32)  # 初始化第二個尺度上的輸出，shape 爲 [4, 26, 26, 3, 85]
            batch_label_lbbox = np.zeros((self.batch_size, self.train_output_sizes[2], self.train_output_sizes[2],
                                          self.anchor_per_scale, 5 + self.num_classes), dtype=np.float32)  # 初始化第三個尺度上的輸出，shape 爲 [4, 13, 13, 3, 85]

            batch_sbboxes = np.zeros((self.batch_size, self.max_bbox_per_scale, 4), dtype=np.float32)
            batch_mbboxes = np.zeros((self.batch_size, self.max_bbox_per_scale, 4), dtype=np.float32)
            batch_lbboxes = np.zeros((self.batch_size, self.max_bbox_per_scale, 4), dtype=np.float32)

            num = 0  # 記錄現在遍歷到一個批次（4張圖片）中的第幾張圖片了
            if self.batch_count < self.num_batchs:  # 如果記錄的 batch 個數還沒達到總 batch 個數
                while num < self.batch_size:  # 如果這個批次（4張圖片）還沒遍歷完
                    index = self.batch_count * self.batch_size + num  # 記錄現在是第幾個樣本（即記錄已經遍歷了的樣本數）
                    if index >= self.num_samples: index -= self.num_samples  # 要是已遍歷樣本數大於總樣本數，將 index 置零，相當於 repeat 操作
                    annotation = self.annotations[index]  # 加載第 index 個樣本的信息（因爲之前有 shuffle 操作，所以這個樣本並非圖片 index.jpg）
                    image, bboxes = self.parse_annotation(annotation)  # 加載圖片及其真實框 -> 數據增強操作 -> 將圖像處理成模型需要輸入的格式
                    label_sbbox, label_mbbox, label_lbbox, sbboxes, mbboxes, lbboxes = self.preprocess_true_boxes(bboxes)  # 把符合要求的三個尺度上的先驗框信息及位置提取出來

                    batch_image[num, :, :, :] = image  # 把批次大小張圖片一起放入數組
                    batch_label_sbbox[num, :, :, :, :] = label_sbbox  # 把批次大小張圖片的在第一尺度上符合要求的先驗框信息一起放入數組
                    batch_label_mbbox[num, :, :, :, :] = label_mbbox  # 把批次大小張圖片的在第二尺度上符合要求的先驗框信息一起放入數組
                    batch_label_lbbox[num, :, :, :, :] = label_lbbox  # 把批次大小張圖片的在第三尺度上符合要求的先驗框信息一起放入數組
                    batch_sbboxes[num, :, :] = sbboxes  # 把批次大小張圖片的在第一尺度上符合要求的先驗框位置一起放入數組
                    batch_mbboxes[num, :, :] = mbboxes  # 把批次大小張圖片的在第二尺度上符合要求的先驗框位置一起放入數組
                    batch_lbboxes[num, :, :] = lbboxes  # 把批次大小張圖片的在第三尺度上符合要求的先驗框位置一起放入數組
                    num += 1
                self.batch_count += 1
                batch_smaller_target = batch_label_sbbox, batch_sbboxes  # 第一尺度目標
                batch_medium_target  = batch_label_mbbox, batch_mbboxes  # 第二尺度目標
                batch_larger_target  = batch_label_lbbox, batch_lbboxes  # 第三尺度目標

                return batch_image, (batch_smaller_target, batch_medium_target, batch_larger_target)
            else:  # 如果記錄的 batch 個數達到總 batch 個數了
                self.batch_count = 0  # 計數歸零
                np.random.shuffle(self.annotations)  # 再次打亂樣本
                raise StopIteration

    def random_horizontal_flip(self, image, bboxes):  # 隨機水平翻轉圖片

        if random.random() < 0.5:  # 執行水平翻轉操作的概率爲 0.5
            _, w, _ = image.shape
            image = image[:, ::-1, :]
            bboxes[:, [0,2]] = w - bboxes[:, [2,0]]

        return image, bboxes

    def random_crop(self, image, bboxes):  # 隨機剪裁圖片

        if random.random() < 0.5:  # 執行剪裁圖片操作的概率爲 0.5
            h, w, _ = image.shape  # 圖片的高和寬
            # 取所有真實框的 [xmin, ymin] 中最小的 [xmin, ymin] 和 [xmax, ymax] 中最大的 [xmax, ymax]
            # 所以 max_bbox 現在所對應的框是面積最大的框
            max_bbox = np.concatenate([np.min(bboxes[:, 0:2], axis=0), np.max(bboxes[:, 2:4], axis=0)], axis=-1)

            max_l_trans = max_bbox[0]  # 最大框離圖片左邊緣的距離
            max_u_trans = max_bbox[1]  # 最大框離圖片上邊緣的距離
            max_r_trans = w - max_bbox[2]  # 最大框離圖片右邊緣的距離
            max_d_trans = h - max_bbox[3]  # 最大框離圖片下邊緣的距離

            crop_xmin = max(0, int(max_bbox[0] - random.uniform(0, max_l_trans)))  # 剪裁後的圖片左上角在原始圖像中的橫座標
            crop_ymin = max(0, int(max_bbox[1] - random.uniform(0, max_u_trans)))  # 剪裁後的圖片左上角在原始圖像中的縱座標
            crop_xmax = max(w, int(max_bbox[2] + random.uniform(0, max_r_trans)))  # 剪裁後的圖片右下角在原始圖像中的橫座標
            crop_ymax = max(h, int(max_bbox[3] + random.uniform(0, max_d_trans)))  # 剪裁後的圖片右下角在原始圖像中的縱座標

            image = image[crop_ymin : crop_ymax, crop_xmin : crop_xmax]  # 剪裁後的圖片

            bboxes[:, [0, 2]] = bboxes[:, [0, 2]] - crop_xmin  # 剪裁後的真實框左上角在原始圖像中的座標
            bboxes[:, [1, 3]] = bboxes[:, [1, 3]] - crop_ymin  # 剪裁後的真實框右下角在原始圖像中的座標

        return image, bboxes

    def random_translate(self, image, bboxes):  # 隨機平移圖片

        if random.random() < 0.5:  # 執行平移圖片操作的概率爲 0.5
            h, w, _ = image.shape  # 圖片的高和寬
            # 取所有真實框的 [xmin, ymin] 中最小的 [xmin, ymin] 和 [xmax, ymax] 中最大的 [xmax, ymax]
            # 所以 max_bbox 現在所對應的框是面積最大的框
            max_bbox = np.concatenate([np.min(bboxes[:, 0:2], axis=0), np.max(bboxes[:, 2:4], axis=0)], axis=-1)

            max_l_trans = max_bbox[0]  # 最大框離圖片左邊緣的距離
            max_u_trans = max_bbox[1]  # 最大框離圖片上邊緣的距離
            max_r_trans = w - max_bbox[2]  # 最大框離圖片右邊緣的距離
            max_d_trans = h - max_bbox[3]  # 最大框離圖片下邊緣的距離

            tx = random.uniform(-(max_l_trans - 1), (max_r_trans - 1))  # 左右平移的距離
            ty = random.uniform(-(max_u_trans - 1), (max_d_trans - 1))  # 上下平移的距離

            M = np.array([[1, 0, tx], [0, 1, ty]])
            image = cv2.warpAffine(image, M, (w, h))  # 仿射變換

            bboxes[:, [0, 2]] = bboxes[:, [0, 2]] + tx  # 平移後的真實框左上角在原始圖像中的座標
            bboxes[:, [1, 3]] = bboxes[:, [1, 3]] + ty  # 平移後的真實框右下角在原始圖像中的座標

        return image, bboxes

    def parse_annotation(self, annotation):

        line = annotation.split()  # 將樣本信息按 ' ' 劃分爲一個列表中的不同元素
        image_path = line[0]  # 提取樣本路徑
        if not os.path.exists(image_path):  # 要是這個路徑不存在就報錯
            raise KeyError("%s does not exist ... " %image_path)
        image = cv2.imread(image_path)  # 加載路徑下對應的圖片
        # 用 map 函數將真實框信息從字符串類型變爲數值類型，然後放到一個列表中（這裏可以參考 text.py 中的說明）
        bboxes = np.array([list(map(int, box.split(','))) for box in line[1:]])

        if self.data_aug:  # 是否對圖片進行數據增強操作
            image, bboxes = self.random_horizontal_flip(np.copy(image), np.copy(bboxes))  # 隨機水平翻轉圖片
            image, bboxes = self.random_crop(np.copy(image), np.copy(bboxes))  # 隨機剪裁圖片
            image, bboxes = self.random_translate(np.copy(image), np.copy(bboxes))  # 隨機平移圖片

        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # BGR --> RGB
        # 將圖像處理成模型需要輸入的格式
        image, bboxes = utils.image_preporcess(np.copy(image), [self.train_input_size, self.train_input_size], np.copy(bboxes))
        return image, bboxes
    
    # 計算兩個框之間的 IOU 值
    def bbox_iou(self, boxes1, boxes2):

        boxes1 = np.array(boxes1)  # 第一個檢測框的座標數據
        boxes2 = np.array(boxes2)  # 第二個檢測框的座標數據

        boxes1_area = boxes1[..., 2] * boxes1[..., 3]  # 第一個檢測框的面積
        boxes2_area = boxes2[..., 2] * boxes2[..., 3]  # 第二個檢測框的面積

        boxes1 = np.concatenate([boxes1[..., :2] - boxes1[..., 2:] * 0.5,
                                boxes1[..., :2] + boxes1[..., 2:] * 0.5], axis=-1)  # 第一個檢測框的左上角座標+右下角座標
        boxes2 = np.concatenate([boxes2[..., :2] - boxes2[..., 2:] * 0.5,
                                boxes2[..., :2] + boxes2[..., 2:] * 0.5], axis=-1)  # 第二個檢測框的左上角座標+右下角座標

        left_up = np.maximum(boxes1[..., :2], boxes2[..., :2])  # 對上圖來說，left_up=[xmin2, ymin2]
        right_down = np.minimum(boxes1[..., 2:], boxes2[..., 2:])  # 對上圖來說，right_down=[xmax1, ymax1]

        inter_section = np.maximum(right_down - left_up, 0.0)  # 交集區域
        inter_area = inter_section[..., 0] * inter_section[..., 1]  # 交集面積
        union_area = boxes1_area + boxes2_area - inter_area  # 並集面積

        return inter_area / union_area

    def preprocess_true_boxes(self, bboxes):

        # label 中有3個列表，每個列表形狀爲 [當前 feature map 尺寸, 當前 feature map 尺寸, 當前 feature map 下的先驗框數量, 5 + 種類數量]
        label = [np.zeros((self.train_output_sizes[i], self.train_output_sizes[i], self.anchor_per_scale,
                           5 + self.num_classes)) for i in range(3)]
        # bboxes_xywh 中有3個列表，每個列表形狀爲 [每個尺度下允許的最大檢測物數量, 檢測框 4 個座標]
        bboxes_xywh = [np.zeros((self.max_bbox_per_scale, 4)) for _ in range(3)]
        bbox_count = np.zeros((3,))  # 記錄最終在各個尺度下有多少符合條件的先驗框

        for bbox in bboxes:  # 對這張圖片上所有真實框中的一個真實框操作
            bbox_coor = bbox[:4]  # 記錄這個真實框的座標（左上角 + 右下角）
            bbox_class_ind = bbox[4]  # 記錄這個真實框中內容的種類（用索引表示）

            onehot = np.zeros(self.num_classes, dtype=np.float)  # 初始化關於種類的獨熱編碼
            onehot[bbox_class_ind] = 1.0  # 把對應這個真實框內容種類的獨熱編碼位置置 1
            # uniform_distribution 是一個形狀爲 [種類數,] 的數組，每個元素值都是種類數的倒數
            uniform_distribution = np.full(self.num_classes, 1.0 / self.num_classes)
            deta = 0.01
            smooth_onehot = onehot * (1 - deta) + deta * uniform_distribution  # 標籤平滑

            # 將 [xmin, ymin, xmax, ymax] 轉換爲 [中心橫座標, 中心縱座標, 寬, 高]
            bbox_xywh = np.concatenate([(bbox_coor[2:] + bbox_coor[:2]) * 0.5, bbox_coor[2:] - bbox_coor[:2]], axis=-1)
            # 將真實框的座標轉化到 feature map 上（轉化完之後就不一定都是整數了）
            bbox_xywh_scaled = 1.0 * bbox_xywh[np.newaxis, :] / self.strides[:, np.newaxis]  # bbox_xywh shape: [1, 4]；strides shape: [3, ]; bbox_xywh_scaled shape: [1, 3, 4]

            iou = []
            exist_positive = False  # False 表示還沒有符合條件的正樣本
            for i in range(3):  # 3 個尺度中的某一尺度下
                anchors_xywh = np.zeros((self.anchor_per_scale, 4))  # 初始化先驗框位置，shape = [3, 4]，3 表示一個尺度下有三種先驗框，4 表示是每個先驗框的中心座標和寬高
                anchors_xywh[:, 0:2] = np.floor(bbox_xywh_scaled[i, 0:2]).astype(np.int32) + 0.5  # 該尺度下三個先驗框的中心座標（讓它在真實框中心所在的格子上）
                anchors_xywh[:, 2:4] = self.anchors[i]  # 該尺度下三個先驗框的寬度和高度

                iou_scale = self.bbox_iou(bbox_xywh_scaled[i][np.newaxis, :], anchors_xywh)  # 計算三個先驗框與（該尺度下）真實框的 IOU 值
                iou.append(iou_scale)
                iou_mask = iou_scale > 0.3  # IOU 值大於 0.3 時此先驗框代表的索引置 1

                if np.any(iou_mask):  # 如果該尺度下有符合條件的先驗框
                    xind, yind = np.floor(bbox_xywh_scaled[i, 0:2]).astype(np.int32)  # 該尺度下真實框的中心座標

                    # 表示在以 (yind, xind) 這個格子爲中心的先驗框裏有目標
                    label[i][yind, xind, iou_mask, :] = 0
                    label[i][yind, xind, iou_mask, 0:4] = bbox_xywh  # 將真實框的中心座標和寬高賦給這個符合條件的先驗框
                    label[i][yind, xind, iou_mask, 4:5] = 1.0  # 1 表示這個先驗框裏有目標
                    label[i][yind, xind, iou_mask, 5:] = smooth_onehot  # 將代表種類的獨熱編碼賦給這個符合條件的先驗框

                    bbox_ind = int(bbox_count[i] % self.max_bbox_per_scale)  # 記錄這是第幾個符合條件的先驗框
                    bboxes_xywh[i][bbox_ind, :4] = bbox_xywh  # 記錄所有符合條件的先驗框的中心座標和寬高
                    bbox_count[i] += 1

                    exist_positive = True  # True 表示有符合條件的正樣本了

            if not exist_positive:  # 如果三個尺度下都沒有先驗框符合條件
                best_anchor_ind = np.argmax(np.array(iou).reshape(-1), axis=-1)  # 取 IOU 值最大的那個先驗框
                best_detect = int(best_anchor_ind / self.anchor_per_scale)  # 哪一個尺度
                best_anchor = int(best_anchor_ind % self.anchor_per_scale)  # 這個尺度下的哪個先驗框
                xind, yind = np.floor(bbox_xywh_scaled[best_detect, 0:2]).astype(np.int32)  # 該尺度下真實框的中心座標

                label[best_detect][yind, xind, best_anchor, :] = 0
                label[best_detect][yind, xind, best_anchor, 0:4] = bbox_xywh  # 將真實框的中心座標和寬高賦給這個先驗框
                label[best_detect][yind, xind, best_anchor, 4:5] = 1.0  # 1 表示這個先驗框裏有目標
                label[best_detect][yind, xind, best_anchor, 5:] = smooth_onehot  # 將代表種類的獨熱編碼賦給這個先驗框

                bbox_ind = int(bbox_count[best_detect] % self.max_bbox_per_scale)  # 記錄這是第幾個符合條件的先驗框
                bboxes_xywh[best_detect][bbox_ind, :4] = bbox_xywh  # 記錄所有符合條件的先驗框的中心座標和寬高
                bbox_count[best_detect] += 1
        label_sbbox, label_mbbox, label_lbbox = label  # 將符合條件的先驗框標籤（包括位置、置信度和種類信息）按照尺度分到三個列表裏
        sbboxes, mbboxes, lbboxes = bboxes_xywh  # 將符合條件的先驗框的中心座標和寬高按照尺度分到三個列表裏
        return label_sbbox, label_mbbox, label_lbbox, sbboxes, mbboxes, lbboxes

    def __len__(self):
        return self.num_batchs