ubuntu如何训练测试python_Ubuntu下利用TensorFlow Object Detection API测试训练的模型

import numpy as np

import os

import six.moves.urllib as urllib

import sys

import tarfile

import tensorflow as tf

import zipfile

import cv2

​

from collections import defaultdict

from io import StringIO

from matplotlib import pyplot as plt

from PIL import Image

from object_detection.utils import label_map_util

from object_detection.utils import visualization_utils as vis_util

​

​

MODEL_NAME = 'workspace/training_demo/trained-inference-graphs/output_inference_graph_v1.pb'

​

​

# Path to frozen detection graph. This is the actual model that is used for the object detection.

PATH_TO_CKPT = MODEL_NAME + '/frozen_inference_graph.pb'

​

# List of the strings that is used to add correct label for each box.

PATH_TO_LABELS = 'workspace/training_demo/annotations/label_map.pbtxt'

​

# Number of classes to detect

NUM_CLASSES = 5

​

​

# Load a (frozen) Tensorflow model into memory.

detection_graph = tf.Graph()

with detection_graph.as_default():

od_graph_def = tf.compat.v1.GraphDef()

with tf.io.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 maps map indices to category names, so that when our convolution network predicts `5`, we know that this corresponds to `airplane`. Here we use internal utility functions, but anything that returns a dictionary mapping integers to appropriate string labels would be fine

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)

​

​

​

data_dir = 'workspace/training_demo/test_data'

​

# Detection

with detection_graph.as_default():

with tf.compat.v1.Session(graph=detection_graph) as sess:

for filename in os.listdir(data_dir):

image_np = np.array(Image.open(os.path.join(data_dir, filename)))

# Expand dimensions since the model expects images to have shape: [1, None, None, 3]

image_np_expanded = np.expand_dims(image_np, axis=0)

# Extract image tensor

image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')

# Extract detection boxes

boxes = detection_graph.get_tensor_by_name('detection_boxes:0')

# Extract detection scores

scores = detection_graph.get_tensor_by_name('detection_scores:0')

# Extract detection classes

classes = detection_graph.get_tensor_by_name('detection_classes:0')

# Extract number of detectionsd

num_detections = detection_graph.get_tensor_by_name(

'num_detections:0')

# Actual detection.

(boxes, scores, classes, num_detections) = sess.run(

[boxes, scores, classes, num_detections],

feed_dict={image_tensor: image_np_expanded})

# Visualization of the results of a detection.

vis_util.visualize_boxes_and_labels_on_image_array(

image_np,

np.squeeze(boxes),

np.squeeze(classes).astype(np.int32),

np.squeeze(scores),

category_index,

use_normalized_coordinates=True,

line_thickness=8)

​

# Display output

# cv2.imshow('object detection', cv2.resize(image_np, (800, 600)))

cv2.imwrite(filename, image_np)