SSD进行物体检测

1.案例效果：

2.案例需求：

代码：

对结果进行标记-完整代码：

def tag_picture(self, images_data, outputs):
    """
    显示预测结果到图片中
    :return:
    """
    # 获取每张图片预测结果的值
    for i, img in enumerate(images_data):
        # 解析输出结果,每张图片的标签，置信度和位置
        pre_label = outputs[i][:, 0]
        pre_conf = outputs[i][:, 1]
        pre_xmin = outputs[i][:, 2]
        pre_ymin = outputs[i][:, 3]
        pre_xmax = outputs[i][:, 4]
        pre_ymax = outputs[i][:, 5]
        print("label: {}, probability: {}, pre_xmin: {}, pre_ymin: {}, pre_xmax: {}, pre_ymax: {}".format(pre_label, pre_conf, pre_xmin, pre_ymin, pre_xmax, pre_ymax))
    # 由于检测出来的物体还是很多，所以进行显示过滤
        top_indices = [i for i, conf in enumerate(pre_conf) if conf >= 0.6]
        top_conf = pre_conf[top_indices]
        top_label_indices = pre_label[top_indices].tolist()
        top_xmin = pre_xmin[top_indices]
        top_ymin = pre_ymin[top_indices]
        top_xmax = pre_xmax[top_indices]
        top_ymax = pre_ymax[top_indices]
        # 画图
        colors = plt.cm.hsv(np.linspace(0, 1, 21)).tolist()
        plt.imshow(img / 255.)
        currentAxis = plt.gca()

        for i in range(top_conf.shape[0]):
            xmin = int(round(top_xmin[i] * img.shape[1]))
            ymin = int(round(top_ymin[i] * img.shape[0]))
            xmax = int(round(top_xmax[i] * img.shape[1]))
            ymax = int(round(top_ymax[i] * img.shape[0]))

            # 获取该图片预测概率，名称，定义显示颜色
            score = top_conf[i]
            label = int(top_label_indices[i])
            label_name = self.classes_name[label - 1]
            display_txt = '{:0.2f}, {}'.format(score, label_name)
            coords = (xmin, ymin), xmax - xmin + 1, ymax - ymin + 1
            color = colors[label]
            # 显示方框
            currentAxis.add_patch(plt.Rectangle(*coords, fill=False, edgecolor=color, linewidth=2))
            # 左上角显示概率以及名称
            currentAxis.text(xmin, ymin, display_txt, bbox={'facecolor': color, 'alpha': 0.5})

        plt.show()

完整代码：

from nets.ssd_net import SSD300
from utils.ssd_utils import  BBoxUtility
from tensorflow.keras.preprocessing.image import load_img, img_to_array
import os
from scipy.misc import imread
from keras.applications.imagenet_utils import preprocess_input
import numpy as np
import matplotlib.pyplot as plt
class SSDTest(object):
    def __init__(self):
        # 定义识别类别
        self.classes_name = ['Aeroplane', 'Bicycle', 'Bird', 'Boat', 'Bottle',
                              'Bus', 'Car', 'Cat', 'Chair', 'Cow', 'Diningtable',
                              'Dog', 'Horse', 'Motorbike', 'Person', 'Pottedplant',
                              'Sheep', 'Sofa', 'Train', 'Tvmonitor']
        # 定义模型输入的参数 1 北京
        self.classes_nums = len(self.classes_name) + 1
        self.input_shape = (300, 300, 3)
    def test(self):
        """
        对于输入的图片进行预测物体位置
        :return:
        """
        # SSD300模型输入以及加载参数
        model = SSD300(self.input_shape, num_classes=self.classes_nums)
        model.load_weights("./ckpt/transfer_01-0.64.h5", by_name=True)
        feature = []
        images_data = []
        # 读取多个本地路径测试图片， preprocessing_input以及保存图像像素值（显示需要）
        for path in os.listdir("./images"):
            img_path = os.path.join("./images/", path)
            # 1. 输入到SSD网络中。数组
            image = load_img(img_path, target_size=(self.input_shape[0], self.input_shape[1]))
            image = img_to_array(image)
            feature.append(image)

            # 2。 读取图片二进制数据
            images_data.append(imread(img_path))




        # 模型预测结果，得到7308个priorbox
        # 处理
        inputs = preprocess_input(np.asarray(feature))
        pred = model.predict(inputs)

        # res形状 (2， 7308， 33) 2代表图片数量，7308代表每个图片预测的default boxes数量， 33：4（位置） + 21（预测概率） + 8（其他default boxes参数）


        # 进行非最大抑制算放处理
        bb = BBoxUtility(self.classes_nums)
        res = bb.detection_out(pred)
        # res (200, 6) (200, 6)
        print(res[0].shape, res[1].shape)
        # 还剩200 个候选框， 每个位置类别都是6
        return res, images_data
    def tag_picture(self, images_data, outputs):
        """
        显示预测结果到图片中
        :return:
        """
        # 获取每张图片预测结果的值
        for i, img in enumerate(images_data):
            # 解析输出结果,每张图片的标签，置信度和位置
            pre_label = outputs[i][:, 0]
            pre_conf = outputs[i][:, 1]
            pre_xmin = outputs[i][:, 2]
            pre_ymin = outputs[i][:, 3]
            pre_xmax = outputs[i][:, 4]
            pre_ymax = outputs[i][:, 5]
            print("label: {}, probability: {}, pre_xmin: {}, pre_ymin: {}, pre_xmax: {}, pre_ymax: {}".format(pre_label, pre_conf, pre_xmin, pre_ymin, pre_xmax, pre_ymax))
        # 由于检测出来的物体还是很多，所以进行显示过滤
            top_indices = [i for i, conf in enumerate(pre_conf) if conf >= 0.6]
            top_conf = pre_conf[top_indices]
            top_label_indices = pre_label[top_indices].tolist()
            top_xmin = pre_xmin[top_indices]
            top_ymin = pre_ymin[top_indices]
            top_xmax = pre_xmax[top_indices]
            top_ymax = pre_ymax[top_indices]
            # 画图
            # 定义21中颜色
            colors = plt.cm.hsv(np.linspace(0, 1, 21)).tolist()
            plt.imshow(img / 255.)
            currentAxis = plt.gca()

            for i in range(top_conf.shape[0]):
                xmin = int(round(top_xmin[i] * img.shape[1]))
                ymin = int(round(top_ymin[i] * img.shape[0]))
                xmax = int(round(top_xmax[i] * img.shape[1]))
                ymax = int(round(top_ymax[i] * img.shape[0]))

                # 获取该图片预测概率，名称，定义显示颜色
                score = top_conf[i]
                label = int(top_label_indices[i])
                label_name = self.classes_name[label - 1]
                display_txt = '{:0.2f}, {}'.format(score, label_name)
                coords = (xmin, ymin), xmax - xmin + 1, ymax - ymin + 1
                color = colors[label]
                # 显示方框
                currentAxis.add_patch(plt.Rectangle(*coords, fill=False, edgecolor=color, linewidth=2))
                # 左上角显示概率以及名称
                currentAxis.text(xmin, ymin, display_txt, bbox={'facecolor': color, 'alpha': 0.5})

            plt.show()

        return None
if __name__ == '__main__':
    ssd = SSDTest()
    outputs, images_data = ssd.test()
    # 显示图片
    ssd.tag_picture(images_data, outputs)