[人工智能]深度学习物体检测实战算法4——视频物体检测测试

视频物体检测测试

SSD训练商品数据

主程序

GPU版本运行

import pickle
from utils.detection_generate import Generator
from utils.ssd_utils import BBoxUtility
from nets.ssd_net import SSD300
from utils.ssd_losses import MuliboxLoss
from tensorflow.python import keras


class SSDTrain(object):
    def __init__(self,num_classes=9,input_shape=(300,300,3),epochs=30):
        """
        初始化网络指定一些参数，训练数据类别，图片需要指定模型输入大小，迭代次数
        :param num_classes:
        """
        self.num_classes = num_classes
        self.batch_size = 32
        self.input_shape = input_shape
        self.epochs = epochs
        #指定训练和读取数据的相关参数
        self.gt_path = "./datasets/commodity_gt.pkl"
        self.image_path = "./datasets/commodity/JPEGImages/"
        prior = pickle.load(open("./datasets/prior_boxes_ssd300.pkl","rb"))
        self.bbox_util = BBoxUtility(self.num_classes,prior)
        self.pre_trained = "./ckpt/pre_trained/weights_SSD300.hdf5"
        #初始化模型
        self.model = SSD300(self.input_shape,num_classes=self.num_classes)




    def get_detefction_data(self):
        """
        获取检测的迭代数据
        :return:
        """
        #1.读取标注数据，构造训练图片名字列表，测试图片名字列表
        gt = pickle.load(open(self.gt_path,'rb'))
        print(gt)
        #图片名字列表
        name_keys = sorted(gt.keys())
        number = int(round(0.8*len(name_keys)))
        #获取训练集长度
        train_keys = name_keys[:number]
        val_keys = name_keys[number:]
        #2.通过generator去获取迭代批次数据
        # gt:所有数据的目标值字典
        # path_prefix:图片的路径
        # image_size:转换成的固定图片大小
        # do_crop: 是否裁剪
        bbox_util = 0
        gen = Generator(gt,self.bbox_util,self.batch_size,self.image_path,
                        train_keys,val_keys,(self.input_shape[0],self.input_shape[1]),do_crop=False)
        return gen
    def init_model_param(self):
        """
        初始化网络模型参数，指定微调的时候，训练部分
        :return:
        """
        #1.加载本地预训练号的模型
        self.model.load_weights(self.pre_trained,by_name = True)
        #2.指定模型当中某些结构冻结freeze
        #冻结模型部分为SSD当中的VGG前半部分
        freeze = ['input_1','conv1_1','conv1_2','pool1',
                  'conv2_1','conv2_2','pool2',
                  'conv3_1','conv3_2','conv3_3','pool3']
        for L in self.model.layers:
            if L.name in freeze:
                L.trainable = False
        return None
    def compile(self):
        """编译模型
        SSD网络的损失函数计算Multiboxloss的compute_loss
        """
        #MultiboxLoss:N个类别+1个背景类别、
        self.model.compile(optimizer = keras.optimizers.Adam(),
                           loss = MuliboxLoss(self.num_classes).compute_loss)
 
if __name__ == '__name__':
    ssd = SSDTrain(num_classes=9)
    gen = ssd.get_detefction_data()