Python 加载YOLOV2模型，实现目标检测

        最近在学习YOLO目标检测的知识点，学习github上面的darkflow项目，项目功能多样复杂。如果想利用训练好的模型实现一些小的应用的话，直接用起来就很不方便，所以最近研究了一下项目里关于yolov2部分预测的代码，并按照自己的要求重新写了一个新的小项目，可以实现加载利用darkflow训练完成的模型，进行目标检测的功能。以备后期做目标检测，可以直接用。项目的git地址：https://github.com/xiaoqiang0008/YOLOV2-predict-use-tensorflow/tree/master

        这个只能做预测，不能训练。训练的部分研究研究，弄明白了再自己重新搞一个，代码不自己重新再写写，总觉得用起来不爽！另外这个仅用于学习，考虑实现基本功能。

环境要求：

系统：windows7_X64

Tensorflow 1.4

Opencv3

先看一下目录结构：

从上往下：

build文件夹存放Cython代码编译的中间文件；

image文件夹存放测试的图片文件；

lib文件夹存放Cython代码相关的模块，该部分主要摘自darkflow项目；

pb文件夹存放训练好的模型文件，包括后缀名为 ’.pb’以及 ’.meta’ 两个文件，前者记录网络结构、权重等变量，后者记录网络的一下参数信息，例如具体类别等；

yolo文件夹存放加载模型文件并实现目标检测的相关模块。

 至于pb模型怎么来的，查看darkflow项目的说明，里面有把darknet模型转换为tensorflow的pb模型的指令，搞一搞就可以了。这里也有两个我转换好的模型tiny_yolo,pb、YOLOV2的pb模型文件，可以下载直接使用。

使用方法说明：

开始之前需要先编译Cython代码。打开命令行窗口，进入到根目录下，执行

python setup.py build_ext --inplace

如果执行上面命令报错的话，可能是因为你的电脑没有安装Cython以及与Python版本相对应的VS的C编译器，这个搞起来也真是麻烦，不过网上有很多解决办法的。

执行 -h 或 --help 查看使用方法

python main.py -h 或 python main.py --help

只有-i 或 --image 必须要指定之外，其他都是由默认值的，不更改的话，直接采用默认即可

接下来可以试着用命令行检测一张照片，看看效果！

python main.py -i image/car.jpg

预测完成，会默认在image文件夹下生成一个名为predict.jpg的图片文件。

当然，pb文件夹中的yolo模型可以预测输出80目标，然而有时候，只想检测几类目标，这样可以利用-l指令了。加上-l的话，程序会自动加载yolo/labels.py 文件下的mylabels列表。我们可以修改这个列表的值，实现只输出制定目标类型。下面看看效果！

先后执行

python main.py -i image/car.jpg
python main.py -i image/car.jpg -l

 

 我的yolo/labels.py中mylabels只设置了一个类别“person”,所以第一张图片只显示检测到了两个人，其他车辆等目标没有显示。

 

使用命令行模式实现预测是为了测试方便，如果想直接嵌入到自己的项目里，可以修改main.py文件，像这样子：

from yolo.predict import YOLO_V2_Predict
from yolo.labels import mylabels
import cv2

image = "image/t.jpg"
pbfile = 'pb/yolo.pb'
metafile = 'pb/yolo.meta'
thresh = 0.2
savedir = 'image/predict.jpg'
labels = mylabels

yolo = YOLO_V2_Predict(pbfile, metafile ,thresh=thresh,classes=labels)
img = yolo.predict(image,save=False)[0]
cv2.imshow('predict',img)
cv2.waitKey(0)

至于代码非常简单，就是加载模型、预测、再就是利用darkflow的代码将网络输出的tensor转换为具体的分类信息。

贴一下主要代码，程序里也有详细的注释就不解释了，记录写来以备自己以后查看方便。

import sys
sys.path.append("..")
import tensorflow as tf
import cv2
import  numpy as np
import ast
from skimage import io, transform
from  lib import cy_yolo2_findboxes
import time

#装饰器
def deco(func):
    '''
    说明：
        装饰器
    :param func:
        回调函数
    :return:
    '''
    def wrapper(*args, **kwargs):
        startTime = time.time()
        f=func(*args, **kwargs)
        endTime = time.time()
        msecs = (endTime - startTime)*1000
        print("time is %d ms" %msecs)
        return f
    return wrapper

class YOLO_V2_Predict(object):
    '''
    YOLO_V2 模型加载，及预测
    '''
    
    def __init__(self, pbfile, meta, thresh=None,classes=None):
        '''
        说明：
            初始化函数，加载网络模型，初始化相关参数
        :param pbfile:
            tensorflow .pb文件路径及名称
        :param meta:
            tensorflow .meta文件路径及名称
        :param thresh:
            设置阈值 [0,1]\n
            None: 使用.meta文件中的模型阈值，一般为0.1\n
            0~1之间的数: 使用设置的阈值
        :param classes:
            设置分类种类，该设置仅仅设置图片显示时是否用矩形框框出指定类别的物体

        '''
        super(YOLO_V2_Predict,self).__init__()
        
        self.meta = self.get_meta(meta)
        if type(thresh) != type(None):
            self.meta['thresh'] = thresh
        
        if type(classes) != type(None) and type(classes) == type([]):
            self.classes = classes
        else:
            self.classes = self.meta['labels']
        
        with tf.Graph().as_default():
            self.output_graph_def = tf.GraphDef()
    
            with open(pbfile, "rb") as f:
                self.output_graph_def.ParseFromString(f.read())
                _ = tf.import_graph_def(self.output_graph_def, name="")
    
            self.sess = tf.Session()
            init = tf.global_variables_initializer()
            self.sess.run(init)
    
            self.input_x = self.sess.graph.get_tensor_by_name("input:0")
            self.inW = int(self.input_x.shape[1])
            self.inH = int(self.input_x.shape[2])
            
            print(self.input_x)
    
            self.out_label = self.sess.graph.get_tensor_by_name("output:0")
            self.outW = int(self.out_label.shape[1])
            self.outH = int(self.out_label.shape[2])
            self.outC = int(self.out_label.shape[3])
            
            print(self.out_label)
        
    def __del__(self):
        print('del sess')
        self.sess.close()

    def get_meta(self,meta):
        '''
        说明：
            加载.meta文件
        :param meta:
            tensorflow .meta文件路径及名称
        :return:
            字典类型，包含网络设置相关参数
        '''
        
        with open(meta,'r') as f:
            data = f.read()
        return ast.literal_eval(data)
    
    def findboxes(self, net_out):
        '''
        说明：
            调用Cython程序，处理网络输出的tensor，计算得到boxes
        :param net_out:
            网络输出的tensor
        :return:
            boxes
        '''
        # meta
        meta = self.meta
        boxes = list()
        boxes=cy_yolo2_findboxes.box_constructor(meta,net_out)
        return boxes
    
    @deco
    def get_net(self,imname):
        '''
        说明：
            读取图片，网络前向计算，得到网络输出
        :param imname:
            图片名称及路径
        :return:
            输出的tensor
        '''
        
        img = io.imread(imname)
        img = transform.resize(img, (self.inW, self.inH, 3))
#        img = cv2.imread(imname)
#        img = cv2.resize(img,(self.inW, self.inH),interpolation=cv2.INTER_AREA)
        img_out_softmax = self.sess.run(self.out_label, feed_dict={self.input_x:np.reshape(img, [-1, self.inW, self.inH, 3])})

        return img_out_softmax.reshape(self.outW,self.outH,self.outC)

    def process_box(self, b, h, w, threshold):
        '''
        说明：
            进一步处理boxes，得到矩形框的详细信息\n
            包括 left, right, top, bot, mess, max_indx, max_prob\n
                  x      y     h    w   类别   颜色      概率值

        :param b:
            boxes
        :param h:
            图片高度
        :param w:
            图片宽度
        :param threshold:
            阈值
        :return:
            预测结果的详细信息\n
            left, right, top, bot, mess, max_indx, max_prob\n
              x      y     h    w   类别   颜色      概率值

        '''
        max_indx = np.argmax(b.probs)
        max_prob = b.probs[max_indx]
        label = self.meta['labels'][max_indx]
        if max_prob > threshold:
            left = int((b.x - b.w / 2.) * w)
            right = int((b.x + b.w / 2.) * w)
            top = int((b.y - b.h / 2.) * h)
            bot = int((b.y + b.h / 2.) * h)
            if left < 0:  left = 0
            if right > w - 1: right = w - 1
            if top < 0:   top = 0
            if bot > h - 1:   bot = h - 1
            mess = '{}'.format(label)
            return (left, right, top, bot, mess, max_indx, max_prob)
        return None

    def postprocess(self, net_out, im, save=True,save_dir = 'image/aaa.jpg'):
        '''
        说明：
            前向网络计算完成后，对网络输出tensor进行处理
        :param net_out:
            前向网络输出的tensor
        :param im:
            待预测的图片
        :param save:
            设置是否保存预测完成的图片
        :param save_dir:
            指定图片保存路径及名称，仅save=True 时有效
        :return:
            预测结果\n
            image:
                预测完成的图片数据
            results:
                以列表形式返回预测结果，列表元素为字典，每个字典中包含每个预测结果的\n
                类别、概率值、包围框四个顶点的坐标值

        '''
        boxes = self.findboxes(net_out)

        # meta
        meta = self.meta
        threshold = meta['thresh']
        colors = meta['colors']
        if type(im) is not np.ndarray:
            imgcv = cv2.imread(im)
        else:
            imgcv = im
        #imgcv = cv2.resize(imgcv,(self.inW, self.inH),interpolation=cv2.INTER_AREA)
        h, w, _ = imgcv.shape

        results = []
        for b in boxes:
            boxResults = self.process_box(b, h, w, threshold)
            if boxResults is None:
                continue
            left, right, top, bot, mess, max_indx, confidence = boxResults
            thick = int((h + w) // 300)
            results.append(
                {"label": mess, "confidence": float('%.2f' % confidence), "topleft": {"x": left, "y": top},
                 "bottomright": {"x": right, "y": bot}})
            if mess in self.classes:
                cv2.rectangle(imgcv,
                              (left, top), (right, bot),
                              colors[max_indx], thick)
                cv2.putText(imgcv, mess+' '+str(confidence), (left, top - 12),
                            0, 1e-3 * h, colors[max_indx], thick // 3)

        if save:
            cv2.imwrite(save_dir, imgcv)
        return imgcv,{"label": mess, "confidence": float('%.2f' % confidence),
                      "topleft": {"x": left, "y": top},"bottomright": {"x": right, "y": bot}}

    @deco    
    def predict(self,image,save = True,save_dir = 'image/predict.jpg',thresh=None):
        '''
        说明：
            预测函数
        :param image:
            待预测的图片名称及路径
        :param save:
            设置是否保存预测完成的图片
        :param save_dir:
            指定图片保存路径及名称，仅save=True 时有效
        :param thresh:
            设置阈值 [0,1]\n
            None: 使用.meta文件中的模型阈值，一般为0.1\n
            0~1之间的数: 使用设置的阈值

        :return:
            预测结果\n
            image:
                预测完成的图片数据
            results:
                以列表形式返回预测结果，列表元素为字典，每个字典中包含每个预测结果的\n
                类别、概率值、包围框四个顶点的坐标值
        '''
        if type(thresh) != type(None):
            self.meta['thresh'] = thresh
        return self.postprocess(self.get_net(image),image,save=save,save_dir=save_dir)

 

参考：

https://github.com/thtrieu/darkflow