TensorFlow object detection API以csv格式输出测试集数据（python）

import time
start = time.time()
import numpy as np
import os
import tensorflow as tf
import pandas as pd
import csv
from PIL import Image
from object_detection.utils import label_map_util
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'



# 要改的内容
###############################################
PATH_TO_CKPT = '/home/user/Deep_learning/Tensorflow_obj/models-master/research/object_detection/voc/export/frozen_inference_graph.pb'   # 模型及标签地址
PATH_TO_LABELS = '/home/user/Deep_learning/Tensorflow_obj/models-master/research/object_detection/voc/pascal_label_map.pbtxt'

NUM_CLASSES = 20           # 检测对象个数

PATH_TO_TEST_IMAGES_DIR = '/home/user/Deep_learning/Tensorflow_obj/models-master/research/object_detection/voc/test/'               # 测试图片路径
output_csv_path = ('/home/user/Deep_learning/Tensorflow_obj/models-master/research/object_detection/voc/VOCdevkit/')
confident = 0  # 置信度，即scores>confident的目标才被输出
###############################################
    
# Load a (frozen) Tensorflow model into memory.
detection_graph = tf.Graph()
with detection_graph.as_default():
  od_graph_def = tf.GraphDef()
  with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
    serialized_graph = fid.read()
    od_graph_def.ParseFromString(serialized_graph)
    tf.import_graph_def(od_graph_def, name='')    
    
    
# Loading label map
label_map = label_map_util.load_labelmap(PATH_TO_LABELS)
categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES, use_display_name=True)
category_index = label_map_util.create_category_index(categories)


os.chdir(PATH_TO_TEST_IMAGES_DIR)
TEST_IMAGE_PATHS = os.listdir(PATH_TO_TEST_IMAGES_DIR)

config = tf.ConfigProto()
config.gpu_options.allow_growth = True
with detection_graph.as_default():
  with tf.Session(graph=detection_graph, config=config) as sess:
    image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
    detection_boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
    detection_scores = detection_graph.get_tensor_by_name('detection_scores:0')
    detection_classes = detection_graph.get_tensor_by_name('detection_classes:0')
    num_detections = detection_graph.get_tensor_by_name('num_detections:0')
    data = pd.DataFrame()
    e=1
    for image_path in TEST_IMAGE_PATHS:     # 开始检测
        image = Image.open(image_path)          # 读图片
        width, height = image.size
        image_np = np.array(image)
        image_np_expanded = np.expand_dims(image_np, axis=0)
        # Actual detection.
        (boxes, scores, classes, num) = sess.run(
            [detection_boxes, detection_scores, detection_classes, num_detections],
            feed_dict={image_tensor: image_np_expanded})

        s_boxes = boxes[scores >= confident]
        s_classes = classes[scores >= confident]
        s_scores = scores[scores >= confident]
        print(e)
        e = e+1
        for i in range(len(s_classes)):
            name = image_path.split("\\")[-1]
            # name = image_path.split("\\")[-1].split('.')[0]   # 不带后缀
            ymin = s_boxes[i][0] * height  # ymin
            xmin = s_boxes[i][1] * width  # xmin
            ymax = s_boxes[i][2] * height  # ymax
            xmax = s_boxes[i][3] * width  # xmax
            score = s_scores[i]
            if s_classes[i] in category_index.keys():
                class_name = category_index[s_classes[i]]['name']  # 得到英文class名称
           # if score==0:
              #  class_name = 0 
           #print("name":", name)
           # print("ymin:", ymin)
           # print("xmin:", xmin)
           # print("ymax:", ymax)
           # print("xmax:", xmax)
           # print("score:", score)
           # print("class:", class_name)
           # print("################")
           # print(i)
            newdata= pd.DataFrame(0, index=range(1), columns=['filename','X1','Y1','X2','Y2','X3','Y3','X4','Y4','type'])
            newdata.iloc[0,0] = name
            newdata.iloc[0,1] = xmin  #x1
            newdata.iloc[0,2] = ymin  #y1
            newdata.iloc[0,3] = xmax    #x2
            newdata.iloc[0,4] = ymin     #y2
            newdata.iloc[0,5] = xmax  #x3
            newdata.iloc[0,6] = ymax    #y3
            newdata.iloc[0,7] = xmin    #x4
            newdata.iloc[0,8] = ymax    #y4
            newdata.iloc[0,9] = class_name
           
            data = data.append(newdata)
        data.to_csv(output_csv_path+'predict'+'.csv',index = False)      

end = time.time()
print("Execution Time: ", end - start)

输出结果样式为：

Excel打开

文本打开

可以根据输出需求自行更改代码