import time
start = time.time()
import numpy as np
import os
import tensorflow as tf
import pandas as pd
import csv
import json
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/广东工业纺丝/guangdong1_round1_testA_20190818/' # 测试图片路径
output_csv_path = ('/home/user/Deep_learning/Tensorflow_obj/models-master/research/object_detection/voc/VOCdevkit/')
confident = 0.1 # 置信度,即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
result=[]
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
xmin = float('%.2f' % xmin)
xmax = float('%.2f' % xmax)
ymin = float('%.2f' % ymin)
ymax = float('%.2f' % ymax)
score = float(s_scores[i])
box = [xmin,ymin,xmax,ymax]
#print(box)
#print(s_boxes)
#bbox.append(xmax)
#bbox.append(ymax)
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)
result.append({'name': name, 'category': int(class_name),'bbox':box,'score': score})
#newdata.iloc[0,0] = name
#newdata.iloc[0,1] = int(class_name)
#newdata.iloc[0,2] = bbox
#newdata.iloc[0,3] = score
fw = open('/home/user/Deep_learning/Tensorflow_obj/models-master/research/object_detection/voc/VOCdevkit/result.json', 'a', encoding='utf8')
json.dump(result, fw, indent=4, separators=(',', ': '))
fw.close()
end = time.time()
print("Execution Time: ", end - start)
输出 json 格式为:
根据自己的需求自行修改代码