Python YOLOv3 加载预训练模型对整个图片文件夹进行预测

浅浅的花香味﹌ 2021-10-30 00:16 433阅读 0赞

加载一次YOLOv3预训练模型即可用于对图片预测，而多次加载预训练模型会导致GPU内存不够用。

2019-08-18 14:18:37.202250: I tensorflow/core/common\_runtime/gpu/gpu\_device.cc:993\] Creating TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 1587 MB memory) -> physical GPU (device: 0, name: GeForce 930M, pci bus id: 0000:04:00.0, compute capability: 5.0)

2019-08-18 14:18:45.700148: I tensorflow/core/common\_runtime/gpu/gpu\_device.cc:993\] Creating TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 107 MB memory) -> physical GPU (device: 0, name: GeForce 930M, pci bus id: 0000:04:00.0, compute capability: 5.0)

需要修改如下的代码（多次加载预训练模型）：

img_list = os.listdir(picture_dir)
    for i in range(0,len(img_list)):
        img_path = os.path.join(picture_dir,img_list[i])
        print(img_path)
        img_ori = cv2.imread(img_path)
        with tf.Session() as sess:
             yolo_model = yolov3(args.num_class, args.anchors)
             with tf.variable_scope('yolov3'):
               pred_feature_maps = yolo_model.forward(input_data, False, 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, args.num_class, max_boxes=200, score_thresh=0.3, nms_thresh=0.45)
             saver = tf.train.Saver()
             saver.restore(sess, args.restore_path)
             boxes_, scores_, labels_ = sess.run([boxes, scores, labels], feed_dict={input_data: img})

wu@wu-X555LF:~/YOLOv3_TensorFlow-master$ python predict1.py 
    2019-08-18 14:18:35.219594: I tensorflow/core/platform/cpu_feature_guard.cc:140] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA
    2019-08-18 14:18:35.330784: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:898] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
    2019-08-18 14:18:35.331240: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1212] Found device 0 with properties: 
    name: GeForce 930M major: 5 minor: 0 memoryClockRate(GHz): 0.941
    pciBusID: 0000:04:00.0
    totalMemory: 1.96GiB freeMemory: 1.82GiB
    2019-08-18 14:18:35.331263: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1312] Adding visible gpu devices: 0
    2019-08-18 14:18:37.202250: I tensorflow/core/common_runtime/gpu/gpu_device.cc:993] Creating TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 1587 MB memory) -> physical GPU (device: 0, name: GeForce 930M, pci bus id: 0000:04:00.0, compute capability: 5.0)
    box coords:
    []
    ******************************
    scores:
    []
    ******************************
    labels:
    []
    
    2019-08-18 14:18:45.699982: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1312] Adding visible gpu devices: 0
    2019-08-18 14:18:45.700148: I tensorflow/core/common_runtime/gpu/gpu_device.cc:993] Creating TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 107 MB memory) -> physical GPU (device: 0, name: GeForce 930M, pci bus id: 0000:04:00.0, compute capability: 5.0)
    
    Traceback (most recent call last):
      File "predict1.py", line 66, in <module>
        pred_feature_maps = yolo_model.forward(input_data, False, False)
      File "/home/wu/YOLOv3_TensorFlow-master/model.py", line 51, in forward
        route_1, route_2, route_3 = darknet53_body(inputs)
      File "/home/wu/YOLOv3_TensorFlow-master/utils/layer_utils.py", line 35, in darknet53_body
        net = conv2d(inputs, 32,  3, strides=1)
      File "/home/wu/YOLOv3_TensorFlow-master/utils/layer_utils.py", line 21, in conv2d
        padding=('SAME' if strides == 1 else 'VALID'))
      File "/usr/local/lib/python2.7/dist-packages/tensorflow/contrib/framework/python/ops/arg_scope.py", line 183, in func_with_args
        return func(*args, **current_args)
      File "/usr/local/lib/python2.7/dist-packages/tensorflow/contrib/layers/python/layers/layers.py", line 1050, in convolution
        outputs = layer.apply(inputs)
      File "/usr/local/lib/python2.7/dist-packages/tensorflow/python/layers/base.py", line 809, in apply
        return self.__call__(inputs, *args, **kwargs)
      File "/usr/local/lib/python2.7/dist-packages/tensorflow/python/layers/base.py", line 680, in __call__
        self.build(input_shapes)
      File "/usr/local/lib/python2.7/dist-packages/tensorflow/python/layers/convolutional.py", line 143, in build
        dtype=self.dtype)
      File "/usr/local/lib/python2.7/dist-packages/tensorflow/python/layers/base.py", line 533, in add_variable
        partitioner=partitioner)
      File "/usr/local/lib/python2.7/dist-packages/tensorflow/python/ops/variable_scope.py", line 1297, in get_variable
        constraint=constraint)
      File "/usr/local/lib/python2.7/dist-packages/tensorflow/python/ops/variable_scope.py", line 1093, in get_variable
        constraint=constraint)
      File "/usr/local/lib/python2.7/dist-packages/tensorflow/python/ops/variable_scope.py", line 431, in get_variable
        return custom_getter(**custom_getter_kwargs)
      File "/usr/local/lib/python2.7/dist-packages/tensorflow/contrib/layers/python/layers/layers.py", line 1613, in layer_variable_getter
        return _model_variable_getter(getter, *args, **kwargs)
      File "/usr/local/lib/python2.7/dist-packages/tensorflow/contrib/layers/python/layers/layers.py", line 1604, in _model_variable_getter
        use_resource=use_resource)
      File "/usr/local/lib/python2.7/dist-packages/tensorflow/contrib/framework/python/ops/arg_scope.py", line 183, in func_with_args
        return func(*args, **current_args)
      File "/usr/local/lib/python2.7/dist-packages/tensorflow/contrib/framework/python/ops/variables.py", line 291, in model_variable
        use_resource=use_resource)
      File "/usr/local/lib/python2.7/dist-packages/tensorflow/contrib/framework/python/ops/arg_scope.py", line 183, in func_with_args
        return func(*args, **current_args)
      File "/usr/local/lib/python2.7/dist-packages/tensorflow/contrib/framework/python/ops/variables.py", line 246, in variable
        use_resource=use_resource)
      File "/usr/local/lib/python2.7/dist-packages/tensorflow/python/ops/variable_scope.py", line 408, in _true_getter
        use_resource=use_resource, constraint=constraint)
      File "/usr/local/lib/python2.7/dist-packages/tensorflow/python/ops/variable_scope.py", line 747, in _get_single_variable
        name, "".join(traceback.format_list(tb))))
    ValueError: Variable yolov3/darknet53_body/Conv/weights already exists, disallowed. Did you mean to set reuse=True or reuse=tf.AUTO_REUSE in VarScope? Originally defined at:
    
      File "/usr/local/lib/python2.7/dist-packages/tensorflow/contrib/framework/python/ops/variables.py", line 246, in variable
        use_resource=use_resource)
      File "/usr/local/lib/python2.7/dist-packages/tensorflow/contrib/framework/python/ops/arg_scope.py", line 183, in func_with_args
        return func(*args, **current_args)
      File "/usr/local/lib/python2.7/dist-packages/tensorflow/contrib/framework/python/ops/variables.py", line 291, in model_variable
        use_resource=use_resource)
    
    wu@wu-X555LF:~/YOLOv3_TensorFlow-master$

修改为（加载一次YOLOv3预训练模型用于对所有图片预测）：

with tf.Session() as sess:
        yolo_model = yolov3(args.num_class, 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, args.num_class, max_boxes=200, score_thresh=0.3, nms_thresh=0.45)
    
        saver = tf.train.Saver()
        saver.restore(sess, args.restore_path)
        
        img_list = os.listdir(picture_dir)
        for i in range(0,len(img_list)):
            img_path = os.path.join(picture_dir,img_list[i])
            img_ori = cv2.imread(img_path) #args.input_image
            if args.letterbox_resize:
                img, resize_ratio, dw, dh = letterbox_resize(img_ori, args.new_size[0], args.new_size[1])
            else:
                height_ori, width_ori = img_ori.shape[:2]
            boxes_, scores_, labels_ = sess.run([boxes, scores, labels], feed_dict={input_data: img})

# coding: utf-8
    
    from __future__ import division, print_function
    
    import tensorflow as tf
    import numpy as np
    import argparse
    import cv2
    import os
    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
    from utils.misc_utils import parse_anchors, read_class_names
    from utils.nms_utils import gpu_nms
    from utils.plot_utils import get_color_table, plot_one_box
    from utils.data_aug import letterbox_resize
    
    from model import yolov3
    
    picture_dir = './data/demo_data'
    output_dir = './predict.txt'
    output_txt = open(output_dir,"w");
    
    parser = argparse.ArgumentParser(description="YOLO-V3 test single image test procedure.")
    #parser.add_argument("input_image", type=str,
    #                    help="The path of the input image.")
    parser.add_argument("--anchor_path", type=str, default="./data/yolo_anchors.txt",
                        help="The path of the anchor txt file.")
    parser.add_argument("--new_size", nargs='*', type=int, default=[416, 416],
                        help="Resize the input image with `new_size`, size format: [width, height]")
    parser.add_argument("--letterbox_resize", type=lambda x: (str(x).lower() == 'true'), default=True,
                        help="Whether to use the letterbox resize.")
    parser.add_argument("--class_name_path", type=str, default="./data/my_data/data.names",
                        help="The path of the class names.")
    parser.add_argument("--restore_path", type=str, default="./data/darknet_weights/best/best_model_Epoch_98_step_50291_mAP_0.7358_loss_3.6063_lr_1e-05",
                        help="The path of the weights to restore.")
    args = parser.parse_args()
    
    args.anchors = parse_anchors(args.anchor_path)
    args.classes = read_class_names(args.class_name_path)
    args.num_class = len(args.classes)
    
    color_table = get_color_table(args.num_class)
    
    with tf.Session() as sess:
        input_data = tf.placeholder(tf.float32, [1, args.new_size[1], args.new_size[0], 3], name='input_data')
        yolo_model = yolov3(args.num_class, 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, args.num_class, max_boxes=200, score_thresh=0.3, nms_thresh=0.45)
    
        saver = tf.train.Saver()
        saver.restore(sess, args.restore_path)
        
        img_list = os.listdir(picture_dir)
        for i in range(0,len(img_list)):
            img_path = os.path.join(picture_dir,img_list[i])
            img_ori = cv2.imread(img_path) #args.input_image
            if args.letterbox_resize:
                img, resize_ratio, dw, dh = letterbox_resize(img_ori, args.new_size[0], args.new_size[1])
            else:
                height_ori, width_ori = img_ori.shape[:2]
            img = cv2.resize(img_ori, tuple(args.new_size))
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            img = np.asarray(img, np.float32)
            img = img[np.newaxis, :] / 255.
            boxes_, scores_, labels_ = sess.run([boxes, scores, labels], feed_dict={input_data: img})
    
            # rescale the coordinates to the original image
            if args.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(args.new_size[0]))
                boxes_[:, [1, 3]] *= (height_ori/float(args.new_size[1]))
    
            print("box coords:")
            print(boxes_)
            print('*' * 30)
            print("scores:")
            print(scores_)
            print('*' * 30)
            print("labels:")
            print(labels_)
            output_txt.write(img_path)  #ig_path
            for j in range(len(boxes_)):
                label_class = ' '+ str(labels_[j]) + ' '
                x0, y0, x1, y1 = boxes_[j]
                locate = str(x0)+' '+str(y0)+' '+str(x1)+' '+str(y1)
                output_txt.write(label_class)
                output_txt.write(locate)
            output_txt.write('\n')
                #print(type(labels_[j])) <type 'numpy.int32'>
                #print(type(boxes_[j]))  <type 'numpy.ndarray'>           
                #output_txt.write(labels_[i])  #classes_label
                #output_txt.write(boxes_[j])  #[x0, y0, x1, y1]
                #plot_one_box(img_ori, [x0, y0, x1, y1], label=args.classes[labels_[i]] + ', {:.2f}%'.format(scores_[i] * 100), color=color_table[labels_[i]])
            #cv2.imshow('Detection result', img_ori)
            #cv2.imwrite('detection_result.jpg', img_ori)
            #cv2.waitKey(0)

预测结果：

wu@wu-X555LF:~/YOLOv3_TensorFlow-master$ python predict_add.py 
    2019-08-18 14:33:33.506594: I tensorflow/core/platform/cpu_feature_guard.cc:140] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA
    2019-08-18 14:33:33.583756: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:898] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
    2019-08-18 14:33:33.584216: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1212] Found device 0 with properties: 
    name: GeForce 930M major: 5 minor: 0 memoryClockRate(GHz): 0.941
    pciBusID: 0000:04:00.0
    totalMemory: 1.96GiB freeMemory: 1.84GiB
    2019-08-18 14:33:33.584248: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1312] Adding visible gpu devices: 0
    2019-08-18 14:33:34.217678: I tensorflow/core/common_runtime/gpu/gpu_device.cc:993] Creating TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 1606 MB memory) -> physical GPU (device: 0, name: GeForce 930M, pci bus id: 0000:04:00.0, compute capability: 5.0)
    box coords:
    [[-195.83458   163.17838   884.7618   1150.6211  ]
     [ -31.841726  391.56403   702.5625    912.8325  ]]
    ******************************
    scores:
    [0.9449485 0.7733519]
    ******************************
    labels:
    [2 2]
    box coords:
    []
    ******************************
    scores:
    []
    ******************************
    labels:
    []
    box coords:
    [[314.3078   468.14548  390.6185   567.1601  ]
     [-41.491894 514.53687  213.59131  609.0764  ]
     [ 15.074331 462.53992  156.97115  728.9914  ]
     [-41.491894 514.53687  213.59131  609.0764  ]
     [ 15.074331 462.53992  156.97115  728.9914  ]
     [229.68753  852.45557  244.8793   954.3424  ]
     [135.6899   759.6658   154.38506  844.6287  ]
     [133.06404  731.6753   143.916    811.71185 ]
     [217.71773  875.35455  236.87387  987.4167  ]]
    ******************************
    scores:
    [0.3414791  0.45428857 0.31732422 0.4882345  0.3644903  0.5662823
     0.46963194 0.32721734 0.3194868 ]
    ******************************
    labels:
    [1 4 4 7 7 9 9 9 9]
    box coords:
    [[110.1114    56.36794  241.52951  316.10345 ]
     [241.23412   49.088387 326.56253  264.94296 ]
     [131.32816   86.27491  210.68886  273.31342 ]]
    ******************************
    scores:
    [0.68328035 0.44021612 0.4340352 ]
    ******************************
    labels:
    [0 0 0]
    box coords:
    []
    ******************************
    scores:
    []
    ******************************
    labels:
    []
    box coords:
    [[265.07657   77.91465  331.14487  305.70837 ]
     [412.98926   76.29398  464.5738   350.3297  ]
     [ 37.823307  31.951302 265.8799   321.9834  ]
     [127.45379   25.41013  214.79816  325.07767 ]
     [387.29068   72.292145 440.42075  344.67874 ]]
    ******************************
    scores:
    [0.82978487 0.64961237 0.5234746  0.44627655 0.3882754 ]
    ******************************
    labels:
    [0 0 0 0 0]
    wu@wu-X555LF:~/YOLOv3_TensorFlow-master$