Faster-RCNN（keras）实现宠物狗识别-标注框的生成

前言：关于使用已经训练好的模型进行标注框的生成，知乎上的那篇文章讲的很详细，在此做一下使用记录。我在GitHub上找了一份Yann Henon大神写的源码，我称其为初始版本，初始版本上有train.py，我是使用此版本进行了初步的学习，并使用此版本进行Faster-RCNN的学习，宠物狗的识别也是基于此项目。但是对于我这个新手来说，用这个版本生成标注框很难。我想到之前下载的Yann Henon大神的最新版本，其中使用Jupyter Notebook 写了一个简单的教程，我使用此教程中的训练模型以及@芯尚刃的相关代码实现了标注框的功能。代码会在之后附上。
准备：
初始版本源码地址：https://github.com/yhenon/keras-frcnn
最新版本源码地址https://github.com/fizyr/keras-retinanet.
数据：使用的生成标签后的分类数据
代码：
注：本文中的注释可能不是很规范，具体的话可以自己理解一下。

# 导入模块
import keras
import h5py

# 导入keras_retinanet
from keras_retinanet.models.resnet import custom_objects
from keras_retinanet.utils.image import read_image_bgr, preprocess_image, resize_image
from keras_retinanet.utils.visualization import draw_box, draw_caption
from keras_retinanet.utils.colors import label_color

# 导入各种模块
import matplotlib.pyplot as plt
import cv2
import os
import numpy as np
import time, datetime

# 使用tensorflow作为使用keras的后端支持
import tensorflow as tf

# 使用环境标志来改变要使用的gpu
# os.environ["CUDA_VISIBLE_DEVICES"] = "1"

def get_session():
    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True
    return tf.Session(config=config)


# 设置修改keras的后台tf会话
keras.backend.tensorflow_backend.set_session(get_session())

# 设置你的模型的保存/加载目录
model_path = os.path.join('.', 'snapshots', 'resnet50_coco_best_v2.0.2.hdf5')

# 加载retinanet模型
model = keras.models.load_model(model_path, custom_objects=custom_objects)

# 加载名称映射以达到可视化目的(因为使用的都是狗的标签，这一句好像没用上)
labels_to_names = {0: 'person', 1: 'bicycle', 2: 'car', 3: 'motorcycle', 4: 'airplane', 5: 'bus', 6: 'train', 7: 'truck', 8: 'boat', 9: 'traffic light', 10: 'fire hydrant', 11: 'stop sign', 12: 'parking meter', 13: 'bench', 14: 'bird', 15: 'cat', 16: 'dog', 17: 'horse', 18: 'sheep', 19: 'cow', 20: 'elephant', 21: 'bear', 22: 'zebra', 23: 'giraffe', 24: 'backpack', 25: 'umbrella', 26: 'handbag', 27: 'tie', 28: 'suitcase', 29: 'frisbee', 30: 'skis', 31: 'snowboard', 32: 'sports ball', 33: 'kite', 34: 'baseball bat', 35: 'baseball glove', 36: 'skateboard', 37: 'surfboard', 38: 'tennis racket', 39: 'bottle', 40: 'wine glass', 41: 'cup', 42: 'fork', 43: 'knife', 44: 'spoon', 45: 'bowl', 46: 'banana', 47: 'apple', 48: 'sandwich', 49: 'orange', 50: 'broccoli', 51: 'carrot', 52: 'hot dog', 53: 'pizza', 54: 'donut', 55: 'cake', 56: 'chair', 57: 'couch', 58: 'potted plant', 59: 'bed', 60: 'dining table', 61: 'toilet', 62: 'tv', 63: 'laptop', 64: 'mouse', 65: 'remote', 66: 'keyboard', 67: 'cell phone', 68: 'microwave', 69: 'oven', 70: 'toaster', 71: 'sink', 72: 'refrigerator', 73: 'book', 74: 'clock', 75: 'vase', 76: 'scissors', 77: 'teddy bear', 78: 'hair drier', 79: 'toothbrush'}


def in_out_put(image):
    image = read_image_bgr(image)

    # 得到图像的宽和高（宽和高会在训练数据的时候用上，之后会说）
    (high, wide, _) = image.shape

    # 网络图像预处理
    image = preprocess_image(image)
    image, scale = resize_image(image)

    # 图像处理 
    start = time.time()
    _, _, boxes, nms_classification = model.predict_on_batch(np.expand_dims(image, axis=0))
    print("processing time: ", time.time() - start)

    # 计算预测的标签和分数 
    predicted_labels = np.argmax(nms_classification[0, :, :], axis=1)
    scores = nms_classification[0,np.arange(nms_classification.shape[1]), predicted_labels]
    # 图像缩放校正 
    boxes /= scale

    return scores, predicted_labels, boxes, high, wide

def getRPN():
    """Detect object classes in an image using pre-computed object proposals."""

    CONF_THRESH = 0.5

    train_lable_path = "E:/wangr/ZOUZHEN/WorkSpace/VSCode/MachineLearning/DeepLearning/Pet_Dog_Identify/data_bases/enhance/Images_enhance_train_lable.txt"
    val_lable_path = "E:/wangr/ZOUZHEN/WorkSpace/VSCode/MachineLearning/DeepLearning/Pet_Dog_Identify/data_bases/enhance/Images_enhance_val_lable.txt"
    test_lable_path = "E:/wangr/ZOUZHEN/WorkSpace/VSCode/MachineLearning/DeepLearning/Pet_Dog_Identify/data_bases/enhance/Images_enhance_test_lable.txt"

    train_lable_rpn_path = "E:/wangr/ZOUZHEN/WorkSpace/VSCode/MachineLearning/DeepLearning/Pet_Dog_Identify/data_bases/enhance/train_lable_rpn.txt"
    val_lable_rpn_path = "E:/wangr/ZOUZHEN/WorkSpace/VSCode/MachineLearning/DeepLearning/Pet_Dog_Identify/data_bases/enhance/val_lable_rpn.txt"
    test_lable_rpn_path = "E:/wangr/ZOUZHEN/WorkSpace/VSCode/MachineLearning/DeepLearning/Pet_Dog_Identify/data_bases/enhance/test_lable_rpn.txt"

    error_hd = tf.gfile.FastGFile("E:/wangr/ZOUZHEN/WorkSpace/VSCode/MachineLearning/DeepLearning/Pet_Dog_Identify/data_bases/enhance/not_found_bbox.txt", "w")

    # 得到训练数据的标注框
    print("开始生成训练数据的标注框》》》》")
    oldtime = datetime.datetime.now() # 开始时间
    hd_lable = tf.gfile.FastGFile(train_lable_path, "r")
    hd_rpn = tf.gfile.FastGFile(train_lable_rpn_path, "w")
    for line in hd_lable.readlines():
        line = line.strip("\n")
        lineinfo = line.split(" ")
        scores, predicted_labels, boxes, high, wide = in_out_put(lineinfo[0])
        inds = np.where(scores >= CONF_THRESH)[0]
        if len(inds) == 0:
            error_hd.write(line + "\n")
            print("not found bbox of %s\n" % line)
            continue
        rpn_num = len(inds)
        rpn_num = str(rpn_num)
        bbox = rpn_num + " "
        for i in inds:
            x1, y1, x2, y2 = boxes[0, i, :].astype(int)
            bbox += str(int(x1)) + "," + str(int(y1)) + "," + str(int(x2)) + "," + str(int(y2)) + " "
        line += " " + str(int(wide)) + "," + str(int(high)) + " " + bbox + "\n"
        hd_rpn.write(line)
    hd_lable.close()
    hd_rpn.close()
    error_hd.write("\n\n")
    error_hd.flush()
    newtime_1 = datetime.datetime.now() # 生成完标签的时间
    print("标注框生成完成！！！用时：", newtime_1 - oldtime)

    # 得到验证数据的标注框
    print("开始生成验证数据的标注框》》》》")
    oldtime = datetime.datetime.now() # 开始时间
    hd_lable = tf.gfile.FastGFile(val_lable_path, "r")
    hd_rpn = tf.gfile.FastGFile(val_lable_rpn_path, "w")
    for line in hd_lable.readlines():
        line = line.strip("\n")
        lineinfo = line.split(" ")
        scores, predicted_labels, boxes, high, wide = in_out_put(lineinfo[0])
        inds = np.where(scores >= CONF_THRESH)[0]
        if len(inds) == 0:
            error_hd.write(line + "\n")
            print("not found bbox of %s\n" % line)
            continue
        rpn_num = len(inds)
        rpn_num = str(rpn_num)
        bbox = rpn_num + " "
        for i in inds:
            x1, y1, x2, y2 = boxes[0, i, :].astype(int)
            bbox += str(int(x1)) + "," + str(int(y1)) + "," + str(int(x2)) + "," + str(int(y2)) + " "
        line += " " + str(int(wide)) + "," + str(int(high)) + " " + bbox + "\n"
        hd_rpn.write(line)
    hd_lable.close()
    hd_rpn.close()
    error_hd.write("\n\n")
    error_hd.flush()
    newtime_1 = datetime.datetime.now() # 生成完标签的时间
    print("标注框生成完成！！！用时：", newtime_1 - oldtime)

    # 得到测试数据的标注框
    print("开始生成测试数据的标注框》》》》")
    oldtime = datetime.datetime.now() # 开始时间
    hd_lable = tf.gfile.FastGFile(test_lable_path, "r")
    hd_rpn = tf.gfile.FastGFile(test_lable_rpn_path, "w")
    for line in hd_lable.readlines():
        line = line.strip("\n")
        lineinfo = line.split(" ")
        scores, predicted_labels, boxes, high, wide = in_out_put(lineinfo[0])
        inds = np.where(scores >= CONF_THRESH)[0]
        if len(inds) == 0:
            error_hd.write(line + "\n")
            print("not found bbox of %s\n" % line)
            continue
        rpn_num = len(inds)
        rpn_num = str(rpn_num)
        bbox = rpn_num + " "
        for i in inds:
            x1, y1, x2, y2 = boxes[0, i, :].astype(int)
            bbox += str(int(x1)) + "," + str(int(y1)) + "," + str(int(x2)) + "," + str(int(y2)) + " "
        line += " " + str(int(wide)) + "," + str(int(high)) + " " + bbox + "\n"
        hd_rpn.write(line)
    hd_lable.close()
    hd_rpn.close()
    error_hd.close()
    newtime_1 = datetime.datetime.now() # 生成完标签的时间
    print("标注框生成完成！！！用时：", newtime_1 - oldtime)

if __name__ == '__main__':
    getRPN()

注：对于生成标注框的数据，我们需要和初始版本的训练数据接口进行对接，我会在下一篇中介绍。