Exemplo n.º 1
0
def train():
    # dataset方法
    train_init_op, val_init_op, image_ids, image, y_true = create_iterator()

    # 是否训练placeholders
    is_training = tf.placeholder(tf.bool, name="phase_train")
    pred_boxes_flag = tf.placeholder(tf.float32, [1, None, None])
    pred_scores_flag = tf.placeholder(tf.float32, [1, None, None])

    # gpu nms 操作
    gpu_nms_op = gpu_nms(pred_boxes_flag, pred_scores_flag,
                         train_args.class_num, train_args.nms_topk,
                         train_args.score_threshold, train_args.nms_threshold)

    # 模型加载
    yolo_model = yolov3(train_args.class_num,
                        train_args.anchors,
                        train_args.use_label_smooth,
                        train_args.use_focal_loss,
                        train_args.batch_norm_decay,
                        train_args.weight_decay,
                        use_static_shape=False)
    with tf.variable_scope('yolov3'):
        pred_feature_maps = yolo_model.forward(image, is_training=is_training)
    # 预测值
    y_pred = yolo_model.predict(pred_feature_maps)
    # loss
    loss = yolo_model.compute_loss(pred_feature_maps, y_true)
    l2_loss = tf.losses.get_regularization_loss()

    tf.summary.scalar('train_batch_statistics/total_loss', loss[0])
    tf.summary.scalar('train_batch_statistics/loss_xy', loss[1])
    tf.summary.scalar('train_batch_statistics/loss_wh', loss[2])
    tf.summary.scalar('train_batch_statistics/loss_conf', loss[3])
    tf.summary.scalar('train_batch_statistics/loss_class', loss[4])
    tf.summary.scalar('train_batch_statistics/loss_l2', l2_loss)
    tf.summary.scalar('train_batch_statistics/loss_ratio', l2_loss / loss[0])

    # 加载除去yolov3/yolov3_head下Conv_6、Conv_14、Conv_22
    saver_to_restore = tf.train.Saver(
        var_list=tf.contrib.framework.get_variables_to_restore(
            include=train_args.restore_include,
            exclude=train_args.restore_exclude))
    # 需要更新的变量
    update_vars = tf.contrib.framework.get_variables_to_restore(
        include=train_args.update_part)
    global_step = tf.Variable(float(train_args.global_step),
                              trainable=False,
                              collections=[tf.GraphKeys.LOCAL_VARIABLES])

    # 学习率
    learning_rate = get_learning_rate(global_step)
    tf.summary.scalar('learning_rate', learning_rate)

    # 是否要保存优化器的参数
    if not train_args.save_optimizer:
        saver_to_save = tf.train.Saver()
        saver_best = tf.train.Saver()

    # 优化器
    train_op = build_optimizer(learning_rate, loss, l2_loss, update_vars,
                               global_step)

    if train_args.save_optimizer:
        saver_to_save = tf.train.Saver()
        saver_best = tf.train.Saver()

    with tf.Session() as sess:
        sess.run([
            tf.global_variables_initializer(),
            tf.local_variables_initializer()
        ])
        print('\033[32m----------- Begin resotre weights  -----------')
        saver_to_restore.restore(sess, train_args.restore_path)
        print('\033[32m----------- Finish resotre weights  -----------')
        merged = tf.summary.merge_all()
        writer = tf.summary.FileWriter(train_args.log_dir, sess.graph)

        print('\n\033[32m----------- start to train -----------\n')
        best_mAP = -np.Inf

        for epoch in range(train_args.total_epoches):  # epoch
            print('\033[32m---------epoch:{}---------'.format(epoch))
            sess.run(train_init_op)  # 初始化训练集dataset
            # 初始化五种损失函数
            loss_total, loss_xy, loss_wh, loss_conf, loss_class\
                = AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter()

            for _ in trange(train_args.train_batch_num):  # batch
                # 优化器. summary, 预测值, gt, 损失, global_step, 学习率
                _, __image_ids, summary, __y_pred, __y_true, __loss, __l2_loss, __global_step, __lr = sess.run(
                    [
                        train_op, image_ids, merged, y_pred, y_true, loss,
                        l2_loss, global_step, learning_rate
                    ],
                    feed_dict={is_training: True})
                print(__l2_loss)
                writer.add_summary(summary, global_step=__global_step)

                # 更新误差
                loss_total.update(__loss[0], len(__y_pred[0]))
                loss_xy.update(__loss[1], len(__y_pred[0]))
                loss_wh.update(__loss[2], len(__y_pred[0]))
                loss_conf.update(__loss[3], len(__y_pred[0]))
                loss_class.update(__loss[4], len(__y_pred[0]))

                # 验证
                if __global_step % train_args.train_evaluation_step == 0 and __global_step > 0:
                    # 召回率,精确率
                    recall, precision = evaluate_on_gpu(
                        sess, gpu_nms_op, pred_boxes_flag, pred_scores_flag,
                        __y_pred, __y_true, train_args.class_num,
                        train_args.nms_threshold)

                    info = "epoch:{},global_step:{} | loss_total:{:.2f}, "\
                        .format(epoch, int(__global_step), loss_total.average)
                    info += "xy:{:.2f},wh:{:.2f},conf:{:.2f},class:{:.2f} | "\
                        .format(loss_xy.average, loss_wh.average, loss_conf.average, loss_class.average)
                    info += 'last batch:rec:{:.3f},prec:{:.3f} | lr:{:.5g}'\
                        .format(recall, precision, __lr)
                    print(info)

                    writer.add_summary(make_summary(
                        'evaluation/train_batch_recall', recall),
                                       global_step=__global_step)
                    writer.add_summary(make_summary(
                        'evaluation/train_batch_precision', precision),
                                       global_step=__global_step)

                    if np.isnan(loss_total.average):
                        raise ArithmeticError('梯度爆炸,修改参数后重新训练')

            # 保存模型
            if epoch % train_args.save_epoch == 0 and epoch > 0:
                if loss_total.average <= 2.:
                    print(
                        '\033[32m ----------- Begin sotre weights-----------')
                    print('\033[32m-loss_total.average{}'.format(
                        loss_total.average))
                    saver_to_save.save(
                        sess, train_args.save_dir +
                        'model-epoch_{}_step_{}_loss_{:.4f}_lr_{:.5g}'.format(
                            epoch, int(__global_step), loss_total.average,
                            __lr))
                    print(
                        '\033[32m ----------- Begin sotre weights  -----------'
                    )

            #  验证集评估评估方法
            if epoch % train_args.val_evaluation_epoch == 0 and epoch >= train_args.warm_up_epoch:  # 要过了warm up
                sess.run(val_init_op)

                val_loss_total, val_loss_xy, val_loss_wh, val_loss_conf, val_loss_class = \
                    AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter()

                val_preds = []
                print(
                    '\033[32m -----Begin computing each pred in one epoch of val data-----------'
                )
                for i in trange(train_args.val_img_cnt):  # 在整个验证集上验证
                    __image_ids, __y_pred, __loss = sess.run(
                        [image_ids, y_pred, loss],
                        feed_dict={is_training: False})
                    pred_content = get_preds_gpu(sess, gpu_nms_op,
                                                 pred_boxes_flag,
                                                 pred_scores_flag, __image_ids,
                                                 __y_pred)

                    val_preds.extend(pred_content)
                    # 更新训练集误差
                    val_loss_total.update(__loss[0])
                    val_loss_xy.update(__loss[1])
                    val_loss_wh.update(__loss[2])
                    val_loss_conf.update(__loss[3])
                    val_loss_class.update(__loss[4])
                    if i % 300 == 0:
                        print(i, "--loss-->", __loss)
                print(
                    '\033[32m -----Finish computing each pred in one epoch of val data-----------'
                )
                # 计算验证集mAP
                rec_total, prec_total, ap_total = AverageMeter(), AverageMeter(
                ), AverageMeter()
                gt_dict = parse_gt_rec(train_args.val_file,
                                       train_args.img_size,
                                       train_args.letterbox_resize)

                print('\033[32m -----Begin calculate mAP-------\033[0m')
                info = 'Epoch: {}, global_step: {}, lr: {:.6g} \n'.format(
                    epoch, __global_step, __lr)  # todo
                for j in range(train_args.class_num):
                    npos, nd, rec, prec, ap = voc_eval(
                        gt_dict,
                        val_preds,
                        j,
                        iou_thres=train_args.eval_threshold,
                        use_07_metric=train_args.use_voc_07_metric)
                    info += 'eval: Class {}: Recall: {:.4f}, Precision: {:.4f}, AP: {:.4f}\n'.format(
                        j, rec, prec, ap)
                    rec_total.update(rec, npos)
                    prec_total.update(prec, nd)
                    ap_total.update(ap, 1)

                mAP = ap_total.average
                info += 'eval: Recall: {:.4f}, Precison: {:.4f}, mAP: {:.4f}\n'\
                    .format(rec_total.average, prec_total.average, mAP)
                info += 'eval: loss: total: {:.2f}, xy: {:.2f}, wh: {:.2f}, conf: {:.2f}, class: {:.2f}\n'\
                    .format(val_loss_total.average, val_loss_xy.average,
                            val_loss_wh.average, val_loss_conf.average, val_loss_class.average)
                print(info)
                logging.info(info)
                print('\033[32m -----Begin calculate mAP-------\033[0m')

                if mAP > best_mAP:
                    best_mAP = mAP
                    saver_best.save(
                        sess,
                        train_args.save_dir +
                        'best_model_Epoch_{}_step_{}_mAP_{:.4f}_loss_{:.4f}_lr_{:.7g}'
                        .format(epoch, int(__global_step), best_mAP,
                                val_loss_total.average, __lr)  # todo
                    )
                writer.add_summary(make_summary('evaluation/val_mAP', mAP),
                                   global_step=epoch)
                writer.add_summary(make_summary('evaluation/val_recall',
                                                rec_total.average),
                                   global_step=epoch)
                writer.add_summary(make_summary('evaluation/val_precision',
                                                prec_total.average),
                                   global_step=epoch)
                writer.add_summary(make_summary(
                    'validation_statistics/total_loss',
                    val_loss_total.average),
                                   global_step=epoch)
                writer.add_summary(make_summary(
                    'validation_statistics/loss_xy', val_loss_xy.average),
                                   global_step=epoch)
                writer.add_summary(make_summary(
                    'validation_statistics/loss_wh', val_loss_wh.average),
                                   global_step=epoch)
                writer.add_summary(make_summary(
                    'validation_statistics/loss_conf', val_loss_conf.average),
                                   global_step=epoch)
                writer.add_summary(make_summary(
                    'validation_statistics/loss_class',
                    val_loss_class.average),
                                   global_step=epoch)
Exemplo n.º 2
0
                # we loaded 1 variable
                i += 1
            # we can load weights of conv layer
            shape = var1.shape.as_list()
            num_params = np.prod(shape)

            var_weights = weights[ptr:ptr + num_params].reshape(
                (shape[3], shape[2], shape[0], shape[1]))
            # remember to transpose to column-major
            var_weights = np.transpose(var_weights, (2, 3, 1, 0))
            ptr += num_params
            assign_ops.append(tf.assign(var1, var_weights,
                                        validate_shape=True))
            i += 1
    return assign_ops


model = yolov3(80, anchors)
with tf.Session() as sess:
    inputs = tf.placeholder(tf.float32, [1, img_size, img_size, 3])

    with tf.variable_scope('yolov3'):
        feature_map = model.forward(inputs)

    saver = tf.train.Saver(var_list=tf.global_variables(scope='yolov3'))

    load_ops = load_weights(tf.global_variables(scope='yolov3'), weight_path)
    sess.run(load_ops)
    saver.save(sess, save_path=save_path)
    print('TensorFlow model checkpoint has been saved to {}'.format(save_path))
Exemplo n.º 3
0
def video_detect(input_args):
    vid = cv2.VideoCapture(input_args.input_video)
    video_frame_cnt = int(vid.get(7))
    video_width = int(vid.get(3))
    video_height = int(vid.get(4))
    video_fps = int(vid.get(5))

    fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
    video_writer = cv2.VideoWriter(pred_args.output_video, fourcc, video_fps, (video_width, video_height))

    with tf.Session() as sess:
        input_data = tf.placeholder(tf.float32, [1, pred_args.new_size[1], pred_args.new_size[0], 3], name='input_data')
        yolo_model = yolov3(pred_args.num_class, pred_args.anchors)
        with tf.variable_scope('yolov3'):
            pred_feature_maps = yolo_model.forward(input_data, False)

        pred_boxes, pred_confs, pred_probs = yolo_model.predict(pred_feature_maps)
        pred_scores = pred_confs * pred_probs
        boxes, scores, labels = gpu_nms(
            pred_boxes, pred_scores, pred_args.num_class,
            max_boxes=200, score_thresh=0.3, nms_thresh=0.45
        )
        saver = tf.train.Saver()
        saver.restore(sess, pred_args.weight_path)

        for i in range(video_frame_cnt):
            ret, img_ori = vid.read()
            if input_args.use_letterbox_resize:
                img, resize_ratio, dw, dh = letterbox_resize(img_ori, pred_args.new_size[0], pred_args.new_size[1])
            else:
                height_ori, width_ori = img_ori.shape[:2]
                img = cv2.resize(img_ori, tuple(pred_args.new_size))
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            img = np.asarray(img, np.float32)
            img = img[np.newaxis, :] / 255.

            start_time = time.time()
            boxes_, scores_, labels_ = sess.run([boxes, scores, labels], feed_dict={input_data: img})
            end_time = time.time()

            if input_args.use_letterbox_resize:
                boxes_[:, [0, 2]] = (boxes_[:, [0, 2]] - dw) / resize_ratio
                boxes_[:, [1, 3]] = (boxes_[:, [1, 3]] - dh) / resize_ratio
            else:
                boxes_[:, [0, 2]] *= (width_ori / float(pred_args.new_size[0]))
                boxes_[:, [1, 3]] *= (height_ori / float(pred_args.new_size[1]))

            for i in range(len(boxes_)):
                x0, y0, x1, y1 = boxes_[i]
                plot_one_box(img_ori, [x0, y0, x1, y1],
                             label=pred_args.classes[labels_[i]] + ', {:.2f}%'.format(scores_[i] * 100),
                             color=pred_args.color_table[labels_[i]])
            cv2.putText(
                img_ori, '{:.2f}ms'.format((end_time - start_time) * 1000),
                (40, 40), 0, fontScale=1, color=(0, 255, 0), thickness=2
            )
            cv2.imshow('Detection result', img_ori)
            video_writer.write(img_ori)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

        vid.release()
        video_writer.release()
Exemplo n.º 4
0
def img_detect(input_args):
    """
    图片检测
    :param input_args:
    :return:
    """
    img_ori = cv2.imread(input_args.input_image)  # opencv 打开
    if input_args.use_letterbox_resize:
        img, resize_ratio, dw, dh = letterbox_resize(img_ori, pred_args.new_size[0], pred_args.new_size[1])
    else:
        height_ori, width_ori = img_ori.shape[:2]
        img = cv2.resize(img_ori, tuple(pred_args.new_size))

    # img 转RGB, 转float, 归一化
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = np.asarray(img, np.float32)
    img = img[np.newaxis, :] / 255.

    sess = tf.Session()

    input_data = tf.placeholder(
        tf.float32, [1, pred_args.new_size[1], pred_args.new_size[0], 3], name='input_data'
    )
    with tf.variable_scope('yolov3'):
        yolo_model = yolov3(pred_args.num_class, pred_args.anchors)
        pred_feature_maps = yolo_model.forward(input_data, False)

    pred_boxes, pred_confs, pred_probs = yolo_model.predict(pred_feature_maps)
    pred_scores = pred_confs * pred_probs
    boxes, scores, labels = gpu_nms(
        pred_boxes, pred_scores, pred_args.num_class,
        max_boxes=200, score_thresh=0.3, nms_thresh=0.45)

    saver = tf.train.Saver()
    saver.restore(sess, pred_args.weight_path)

    boxes_, scores_, labels_ = sess.run([boxes, scores, labels], feed_dict={input_data: img})

    # 还原坐标到原图
    if input_args.use_letterbox_resize:
        boxes_[:, [0, 2]] = (boxes_[:, [0, 2]] - dw) / resize_ratio
        boxes_[:, [1, 3]] = (boxes_[:, [1, 3]] - dh) / resize_ratio
    else:
        boxes_[:, [0, 2]] *= (width_ori / float(pred_args.new_size[0]))
        boxes_[:, [1, 3]] *= (height_ori / float(pred_args.new_size[1]))

    print('box coords:', boxes_, '\n' + '*' * 30)
    print('scores:', scores_, '\n' + '*' * 30)
    print('labels:', labels_)

    for i in range(len(boxes_)):
        x0, y0, x1, y1 = boxes_[i]
        plot_one_box(
            img_ori, [x0, y0, x1, y1],
            label=pred_args.classes[labels_[i]] + ', {:.2f}%'.format(scores_[i] * 100),
            color=pred_args.color_table[labels_[i]]
        )
    cv2.imshow('Detection result', img_ori)
    cv2.imwrite(pred_args.output_image, img_ori)
    cv2.waitKey(0)
    sess.close()
Exemplo n.º 5
0
# This script is used to remove the optimizer parameters in the saved checkpoint files.
# These parameters are useless in the forward process.
# Removing them will shrink the checkpoint size a lot.

import sys
sys.path.append('..')

import os
import tensorflow as tf
from net.model import yolov3

# params
ckpt_path = ''
class_num = 20
save_dir = 'shrinked_ckpt'
if not os.path.exists(save_dir):
    os.makedirs(save_dir)

image = tf.placeholder(tf.float32, [1, 416, 416, 3])
yolo_model = yolov3(class_num, None)
with tf.variable_scope('yolov3'):
    pred_feature_maps = yolo_model.forward(image)

saver_to_restore = tf.train.Saver()
saver_to_save = tf.train.Saver()

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    saver_to_restore.restore(sess, ckpt_path)
    saver_to_save.save(sess, save_dir + '/shrinked')
Exemplo n.º 6
0
    def __pre_operate(self):
        """
        初始化部分操作
        :return:
        """
        # gpu nms 操作
        self.gpu_nms_op = gpu_nms(self.pred_boxes_flag, self.pred_scores_flag,
                                  train_args.class_num, train_args.nms_topk,
                                  train_args.score_threshold,
                                  train_args.nms_threshold)

        # 模型加载
        yolo_model = yolov3(train_args.class_num,
                            train_args.anchors,
                            train_args.use_label_smooth,
                            train_args.use_focal_loss,
                            train_args.batch_norm_decay,
                            train_args.weight_decay,
                            use_static_shape=False)

        with tf.variable_scope('yolov3'):
            pred_feature_maps = yolo_model.forward(
                self.image, is_training=self.is_training)

        # 预测值
        self.y_pred = yolo_model.predict(pred_feature_maps)
        # loss
        self.loss = yolo_model.compute_loss(pred_feature_maps, self.y_true)
        self.l2_loss = tf.losses.get_regularization_loss()
        # 学习率
        self.learning_rate = get_learning_rate(self.global_step)
        self.__loss_summary()

        # 加载Saver
        self.saver_to_restore = tf.train.Saver(
            var_list=tf.contrib.framework.get_variables_to_restore(
                include=train_args.restore_include,
                exclude=train_args.restore_exclude))

        # 是否要保存优化器的参数
        if not train_args.save_optimizer:
            self.saver_to_save = tf.train.Saver()
            self.saver_best = tf.train.Saver()

        # 需要更新的变量
        self.update_vars = tf.contrib.framework.get_variables_to_restore(
            include=train_args.update_part)
        # 优化器
        self.train_op = build_optimizer(self.learning_rate, self.loss,
                                        self.l2_loss, self.update_vars,
                                        self.global_step)

        if train_args.save_optimizer:
            self.saver_to_save = tf.train.Saver()
            self.saver_best = tf.train.Saver()

        self.sess.run([
            tf.global_variables_initializer(),
            tf.local_variables_initializer()
        ])
        print('\033[32m----------- Begin resotre weights  -----------\033[0m')
        self.saver_to_restore.restore(self.sess, train_args.restore_path)
        print('\033[32m----------- Finish resotre weights  -----------\033[0m')
        self.merged = tf.summary.merge_all()