文章目录
- 文章说明
- 完整代码
文章说明
本系列文章旨在对 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 cfgclass 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)) # 一共有几个 batchself.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 annotationsdef __iter__(self):return selfdef __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 += 1self.batch_count += 1batch_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 StopIterationdef random_horizontal_flip(self, image, bboxes): # 随机水平翻转图片if random.random() < 0.5: # 执行水平翻转操作的概率为 0.5_, w, _ = image.shapeimage = image[:, ::-1, :]bboxes[:, [0,2]] = w - bboxes[:, [2,0]]return image, bboxesdef random_crop(self, image, bboxes): # 随机剪裁图片if random.random() < 0.5: # 执行剪裁图片操作的概率为 0.5h, 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, bboxesdef random_translate(self, image, bboxes): # 随机平移图片if random.random() < 0.5: # 执行平移图片操作的概率为 0.5h, 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, bboxesdef 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_areadef 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.01smooth_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 时此先验框代表的索引置 1if 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, :] = 0label[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] += 1exist_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, :] = 0label[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] += 1label_sbbox, label_mbbox, label_lbbox = label # 将符合条件的先验框标签（包括位置、置信度和种类信息）按照尺度分到三个列表里sbboxes, mbboxes, lbboxes = bboxes_xywh # 将符合条件的先验框的中心坐标和宽高按照尺度分到三个列表里return label_sbbox, label_mbbox, label_lbbox, sbboxes, mbboxes, lbboxesdef __len__(self):return self.num_batchs