【YOLO系列】YOLOv3代码详解(二)：检测脚本yolo.py

前言

        以下内容仅为个人在学习人工智能中所记录的笔记，先将目标识别算法yolo系列的整理出来分享给大家，供大家学习参考。

        本文仅对YOLOV3代码中关键部分进行了注释，未掌握基础代码的铁汁可以自己百度一下。

        若文中内容有误，希望大家批评指正。

---

资料下载

        YOLOV3论文下载地址：YOLOv3：An Incremental Improvement

回顾

        YOLO V1：【YOLO系列】YOLO V1论文思想详解

        YOLO V2：【YOLO系列】YOLO V2论文思想详解

        YOLO V3：【YOLO系列】 YOLOv3论文思想详解

项目地址

        YOLOV3 keras版本：下载地址

        YOLOV3 Tensorflow版本：下载地址

        YOLOV3 Pytorch版本：下载地址

Gitee仓库

        YOLOV3 各版本：yolov3各版本

---

        本文主要基于keras版本进行讲解

        话不多说，直接上代码

---

一、yolo.py脚本代码详解

        yolo.py脚本主要用于评估输入的图像，输出检测的目标，并在图像中绘制检测出的目标与置信度。

1、设置默认参数

        包括模型文件、Anchor Box、类别文件、检测阈值、IOU阈值、图像大小以及使用的gpu数量

class YOLO(object):
    _defaults = {
        "model_path": 'model_data/yolo.h5',                # 训练好的模型文件路径
        "anchors_path": 'model_data/yolo_anchors.txt',     # 聚类生成的Anchor Box文件路径
        "classes_path": 'model_data/coco_classes.txt',     # coco数据集的类别文件路径
        "score": 0.3,                                      # 目标检测阈值
        "iou": 0.45,                                       # iou阈值
        "model_image_size": (416, 416),                    # 输入图像的大小
        "gpu_num": 1,                                      # 使用的gpu数量
    }

 2、设置classmethod装饰器

        用于外部调用获取相关信息

# 设置classmethod装饰器，用于获取_defaults中的值
    @classmethod
    def get_defaults(cls, n):
        if n in cls._defaults:
            return cls._defaults[n]
        else:
            return "Unrecognized attribute name '" + n + "'"

3、初始化YOLO类参数

        包括class_names，anchor，创建计算图，调用generate()方法获取boxes，score，classes

# 初始化类方法，获取class_names，anchor，session，boxes，score，classes。
    def __init__(self, **kwargs):
        self.__dict__.update(self._defaults)                   # set up default values
        self.__dict__.update(kwargs)                           # and update with user overrides
        self.class_names = self._get_class()                   # 获取类别的名称
        self.anchors = self._get_anchors()                     # 获取Anchors大小
        self.sess = K.get_session()                            # 建立的session计算图
        self.boxes, self.scores, self.classes = self.generate()

4、获取类别与Anchors

    def _get_class(self):
        # os.path.expanduser()用于将路径字符串中的波浪线(~)扩展为用户的主目录，波浪线(~)一般在liunx中较多
        classes_path = os.path.expanduser(self.classes_path)
        with open(classes_path) as f:
            class_names = f.readlines()
        class_names = [c.strip() for c in class_names]
        return class_names

    def _get_anchors(self):
        anchors_path = os.path.expanduser(self.anchors_path)
        with open(anchors_path) as f:
            anchors = f.readline()
        anchors = [float(x) for x in anchors.split(',')]
        return np.array(anchors).reshape(-1, 2)

5、generate()函数输出图片目标框、置信度、类别

        （1）加载训练好的model文件；

        （2）为所有的类别生成一个边框的颜色；

        （3）创建输入图片tensor；

        （4）调用评估函数输出检测图片目标框、置信度、类别。

    def generate(self):
        # 获取model的路径
        model_path = os.path.expanduser(self.model_path)
        # 判断model是否以h5结尾
        assert model_path.endswith('.h5'), 'Keras model or weights must be a .h5 file.'

        # Load model, or construct model and load weights.
        # num_anchors = 9，yolov3有9个先验框
        num_anchors = len(self.anchors)
        # num_classes = 80，coco集一共80类
        num_classes = len(self.class_names)
        # 判断是否为tiny版本，如果是，则加载tiny model
        is_tiny_version = num_anchors == 6  # default setting
        try:
            self.yolo_model = load_model(model_path, compile=False)
        except:
            self.yolo_model = tiny_yolo_body(Input(shape=(None, None, 3)), num_anchors//2, num_classes) \
                if is_tiny_version else yolo_body(Input(shape=(None, None, 3)), num_anchors//3, num_classes)
            self.yolo_model.load_weights(self.model_path)  # make sure model, anchors and classes match
        else:
            # output_shape[-1]:输出维度的最后一维。 -> (?,13,13,255)->255
            # 255 = (9/3)*(80+5). 9/3:每层特征图对应3个anchor box  80:80个类别 5:4+1,框的4个值+1个置信度
            assert self.yolo_model.layers[-1].output_shape[-1] == \
                num_anchors/len(self.yolo_model.output) * (num_classes + 5), \
                'Mismatch between model and given anchor and class sizes'

        print('{} model, anchors, and classes loaded.'.format(model_path))

        # 为所有的类别生成一个边框的颜色。[h,s,v]
        # h(色调）：x/len(self.class_names)  s(饱和度）：1.0  v(明亮）：1.0
        # 对于80种coco目标，确定每一种目标框的绘制颜色，即：将(x/80, 1.0, 1.0)的颜色转换为RGB格式，并随机调整颜色以便于肉眼识别，
        # 其中：一个1.0表示饱和度，一个1.0表示亮度
        hsv_tuples = [(x / len(self.class_names), 1., 1.)
                      for x in range(len(self.class_names))]
        # hsv转换为rgb
        self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
        self.colors = list(
            # hsv取值范围在[0,1]，而RBG取值范围在[0,255]，所以乘上255
            map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
                self.colors))
        # 产生随机种子,固定种子为一致的颜色
        np.random.seed(10101)  # Fixed seed for consistent colors across runs.
        # 打乱，调整颜色，避免相近颜色来装饰相邻的类
        np.random.shuffle(self.colors)  # Shuffle colors to decorrelate adjacent classes.
        # 重置种子为默认
        np.random.seed(None)  # Reset seed to default.

        # Generate output tensor targets for filtered bounding boxes.
        # K.placeholder:keras中的占位符  相当于分配空间
        # 这里是给需要检测的图片预留的，生成一个tensor，输入来自后面detect_image()函数
        self.input_image_shape = K.placeholder(shape=(2, ))
        # 若GPU个数大于等于2，调用multi_gpu_model()
        if self.gpu_num >= 2:
            self.yolo_model = multi_gpu_model(self.yolo_model, gpus=self.gpu_num)
        # yolo_eval(): yolo评估函数
        boxes, scores, classes = yolo_eval(self.yolo_model.output, self.anchors,
                                           len(self.class_names), self.input_image_shape,
                                           score_threshold=self.score, iou_threshold=self.iou)
        return boxes, scores, classes

6、detect_image()函数预测并绘制目标框

        （1）图片尺寸处理：将输入的图片按最长边确定一个缩放比例，然后按比例缩放（采样方法：BICUBIC）图片，再将缩放后的图片粘贴到一个用“绝对灰”R128-G128-B128填充的416x416新图片上，缩放后图片以外的部分保留为灰色；

        （2）归一化图片数值，再添加一个维度生成(bitch, w, h, c)格式，用于model的输入层，调用计算图计算图片目标框、置信度、类别；

        （3）使用Pillow库绘制边框，设置边框宽度，绘制边框和类别字体，将检测出来的所有目标框在图片中绘制出来，输出图片

    def detect_image(self, image):
        start = timer()  # 定时器

        if self.model_image_size != (None, None):
            # 要求进行检测的图片尺寸是32的倍数，因为在Darknet网络中，执行了5次步长为2的卷积操作，即
            # 图片的默认尺寸是416*416，因为在最底层中的特征图大小是13*13，所以13*32=416
            assert self.model_image_size[0] % 32 == 0, 'Multiples of 32 required'
            assert self.model_image_size[1] % 32 == 0, 'Multiples of 32 required'
            # 调用letterbox_image()函数，即：将输入的图片按最长边确定一个比例，然后按比例缩放（采样方法：BICUBIC）图片，
            # 再生成一个用“绝对灰”R128-G128-B128填充的416x416新图片后将缩放后的输入图片粘贴上去，粘贴不到的部分保留为灰色
            boxed_image = letterbox_image(image, tuple(reversed(self.model_image_size)))
        else:
            new_image_size = (image.width - (image.width % 32),
                              image.height - (image.height % 32))
            boxed_image = letterbox_image(image, new_image_size)
        image_data = np.array(boxed_image, dtype='float32')

        print(image_data.shape)  # (416，416,3)
        # 将缩放后图片的数值除以255，做归一化
        image_data /= 255.
        # 在图片前面添加一个维度 -> (1,416,416,3) 满足网络的输入格式 -> (bitch, w, h, c)
        image_data = np.expand_dims(image_data, 0)  # Add batch dimension.

        # 计算boxes,scores,classes，这是使用的是之前建立的session()计算图
        # 即调用generate()函数，将feed_dict中的图像尺寸传递给generate()函数中的placeholder
        # 图片做为model的输入
        out_boxes, out_scores, out_classes = self.sess.run(
            [self.boxes, self.scores, self.classes],
            feed_dict={
                self.yolo_model.input: image_data,  # 图像数据
                self.input_image_shape: [image.size[1], image.size[0]],  # 图像尺寸416x416
                K.learning_phase(): 0  # 学习模式： 0：测试模型；1：训练模式
            })

        print('Found {} boxes for {}'.format(len(out_boxes), 'img'))

        # 使用Pillow库绘制边框，设置边框宽度，绘制边框和类别字体
        # 设置字体
        font = ImageFont.truetype(font='font/FiraMono-Medium.otf',
                    size=np.floor(3e-2 * image.size[1] + 0.5).astype('int32'))
        # 设置目标框线条的宽度
        thickness = (image.size[0] + image.size[1]) // 300

        # 对于c个目标类别中的每个目标框i，调用Pillow画图
        for i, c in reversed(list(enumerate(out_classes))):
            # 目标类别的名字
            predicted_class = self.class_names[c]
            # 框
            box = out_boxes[i]
            # 置信度
            score = out_scores[i]
            # 标签：类别名称+置信度
            label = '{} {:.2f}'.format(predicted_class, score)
            # 加载输入的原始图片
            draw = ImageDraw.Draw(image)
            # 返回标签文字label按照font字体与大小的宽和高（多少个pixels）
            label_size = draw.textsize(label, font)

            top, left, bottom, right = box
            # 目标框的上、左两个坐标小数点后一位向下取整
            top = max(0, np.floor(top + 0.5).astype('int32'))
            left = max(0, np.floor(left + 0.5).astype('int32'))
            # 目标框的下、右两个坐标小数点后一位向下取整，与图片的尺寸相比，取最小值
            bottom = min(image.size[1], np.floor(bottom + 0.5).astype('int32'))
            right = min(image.size[0], np.floor(right + 0.5).astype('int32'))
            print(label, (left, top), (right, bottom))
            # 确定标签（label）起始点位置
            if top - label_size[1] >= 0:
                text_origin = np.array([left, top - label_size[1]])
            else:
                text_origin = np.array([left, top + 1])

            # My kingdom for a good redistributable image drawing library.
            # 绘制目标框，线条宽度为thickness
            for i in range(thickness):
                draw.rectangle(
                    [left + i, top + i, right - i, bottom - i],
                    outline=self.colors[c])
            # 画标签框
            # 绘制一个矩形框，填充颜色作为文字背景
            draw.rectangle(
                [tuple(text_origin), tuple(text_origin + label_size)],
                fill=self.colors[c])
            # 填写标签内容
            draw.text(text_origin, label, fill=(0, 0, 0), font=font)
            del draw
        # 结束计时
        end = timer()
        print(end - start)
        return image

7、detect_video()函数用于视频检测

        （1）打开视频文件，获取视频视频编解码器、视频的帧率、宽度与高度；

        （2）从视频文件中读取每一帧进行检测；

        （3）将文本（FPS）添加到图像（result）上，包括文本内容（text），文本起始位置（org），字体类型（fontFace），字体大小（fontScale），字体颜色（红色），文本线的粗细；

        （4）将检测完成的图片写入out中，生成新的视频。

def detect_video(yolo, video_path, output_path=""):
    import cv2
    # 打开视频
    vid = cv2.VideoCapture(video_path)
    # 判断视频文件是否已成功打开
    if not vid.isOpened():
        raise IOError("Couldn't open webcam or video")
    video_FourCC    = int(vid.get(cv2.CAP_PROP_FOURCC))           # 标识视频编解码器
    video_fps       = vid.get(cv2.CAP_PROP_FPS)                   # 获取视频的帧率
    video_size      = (int(vid.get(cv2.CAP_PROP_FRAME_WIDTH)),
                        int(vid.get(cv2.CAP_PROP_FRAME_HEIGHT)))  # 获取视频的宽度与高度
    isOutput = True if output_path != "" else False
    if isOutput:
        print("!!! TYPE:", type(output_path), type(video_FourCC), type(video_fps), type(video_size))
        out = cv2.VideoWriter(output_path, video_FourCC, video_fps, video_size)        # 创建一个新的视频文件
    accum_time = 0
    curr_fps = 0
    fps = "FPS: ??"
    prev_time = timer()
    while True:
        return_value, frame = vid.read()    # 从视频文件中读取一帧，返回两个元素，第一个为布尔值，判断是否成功读取了帧，第二个元素为读取的帧本身，为一个数组
        image = Image.fromarray(frame)      # 生成图片
        image = yolo.detect_image(image)    # 检测图片
        result = np.asarray(image)
        # 计算当前图片检测时间，累计检测时间
        curr_time = timer()
        exec_time = curr_time - prev_time
        prev_time = curr_time
        accum_time = accum_time + exec_time
        curr_fps = curr_fps + 1
        if accum_time > 1:
            accum_time = accum_time - 1
            fps = "FPS: " + str(curr_fps)
            curr_fps = 0
        # 将文本（FPS）添加到图像（result）上，包括文本内容（text），文本起始位置（org），字体类型（fontFace），字体大小（fontScale），字体颜色（红色），文本线的粗细。
        cv2.putText(result, text=fps, org=(3, 15), fontFace=cv2.FONT_HERSHEY_SIMPLEX,
                    fontScale=0.50, color=(255, 0, 0), thickness=2)
        cv2.namedWindow("result", cv2.WINDOW_NORMAL)
        cv2.imshow("result", result)
        # 将检测完成的图片写入out中，生成新的视频
        if isOutput:
            out.write(result)
        # cv2.waitKey(1) 等待1毫秒键盘输入，返回输入值的ASCII值
        # 0xFF掩码操作，用于确保只获取低 8 位（即一个字节）的数值，
        # cv2.waitKey(1) & 0xFF判断完后，再判断输出结果是否等于 ord('q')
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    yolo.close_session()