物体检测2:物体检测API Guide(TF1)
2020-02-22 本文已影响0人
古风子

物体检测官方demo实例
以下为可在jupyter notebook中运行的代码,完整测试用例,请查看
安装TensorFlow Object Detection API
注意:
因为代码执行会下载模型文件,代码中将下载的代码注释掉了;手动下载好后,放到object_detection目录,并解压到当前目录。
使用以下命令下载:
base) jdf@jdf:~/tensorflow/models/research/object_detection$ wget http://download.tensorflow.org/models/object_detection/ssd_mobilenet_v1_coco_2017_11_17.tar.gz
#####import 各种依赖包
import numpy as np
import os
import six.moves.urllib as urllib
import sys
import tarfile
import tensorflow as tf
import zipfile
import time
import datetime
import matplotlib
matplotlib.use('Agg')
import pylab
from distutils.version import StrictVersion
from collections import defaultdict
from io import StringIO
from PIL import Image
from matplotlib import pyplot as plt
#定义打印函数,记录执行时间
def printf(info):
print(str(datetime.datetime.now())+':'+info)
# This is needed since the notebook is stored in the object_detection folder.
# 防止你不在object_detection目录,因此将当前目录的父目录添加到系统环境变量
sys.path.append("..")
from object_detection.utils import ops as utils_ops
if StrictVersion(tf.__version__) < StrictVersion('1.9.0'):
raise ImportError('Please upgrade your TensorFlow installation to v1.9.* or later!')
##### Env setup
# This is needed to display the images.
# %matplotlib inline
###Object detection imports
#Here are the imports from the object detection module.
from utils import label_map_util
from utils import visualization_utils as vis_util
###Model preparation
# What model to download.
#获取模型文件ssd_mobilenet_v1_coco_2017_11_17.tar.gz下载地址
MODEL_NAME = 'ssd_mobilenet_v1_coco_2017_11_17'
MODEL_FILE = MODEL_NAME + '.tar.gz'
DOWNLOAD_BASE = 'http://download.tensorflow.org/models/object_detection/'
# Path to frozen detection graph. This is the actual model that is used for the object detection.
PATH_TO_FROZEN_GRAPH = MODEL_NAME + '/frozen_inference_graph.pb'
# List of the strings that is used to add correct label for each box.
PATH_TO_LABELS = os.path.join('data', 'mscoco_label_map.pbtxt')
#以下代码是从网址下载模型文件,我这里是在当前object_detection目录使用以下命令下载,并解压到当前目录,因此,将下载代码屏蔽掉了
#wget http://download.tensorflow.org/models/object_detection/ssd_mobilenet_v1_coco_2017_11_17.tar.gz
###Download Model
#opener = urllib.request.URLopener()
#print(DOWNLOAD_BASE + MODEL_FILE)
#opener.retrieve(DOWNLOAD_BASE + MODEL_FILE, MODEL_FILE)
#tar_file = tarfile.open(MODEL_FILE)
#print('tar_file:'+tar_file)
#for file in tar_file.getmembers():
# file_name = os.path.basename(file.name)
# print('file:'+file)
# print('file_name:'+file_name)
# if 'frozen_inference_graph.pb' in file_name:
# tar_file.extract(file, os.getcwd())
###Load a (frozen) Tensorflow model into memory
#加载pb模型文件
printf('start load model file:'+PATH_TO_FROZEN_GRAPH)
detection_graph = tf.Graph()
with detection_graph.as_default():
od_graph_def = tf.GraphDef()
with tf.gfile.GFile(PATH_TO_FROZEN_GRAPH, 'rb') as fid:#object_detection/frozen_inference_graph.pb文件
serialized_graph = fid.read()
od_graph_def.ParseFromString(serialized_graph)
tf.import_graph_def(od_graph_def, name='')
printf('finsh load model file...')
###Loading label map
#PATH_TO_LABELS :#object_detection/data/mscoco_label_map.pbtxt
printf('start load label map file:'+PATH_TO_LABELS)
category_index = label_map_util.create_category_index_from_labelmap(PATH_TO_LABELS, use_display_name=True)
printf('finish load label map file...')
###Helper code
#将图片数据抓换成array数组形式
def load_image_into_numpy_array(image):
(im_width, im_height) = image.size
return np.array(image.getdata()).reshape(
(im_height, im_width, 3)).astype(np.uint8)
printf('finish convert image data to array...')
###Detection
# For the sake of simplicity we will use only 2 images:
# image1.jpg
# image2.jpg
# If you want to test the code with your images, just add path to the images to the TEST_IMAGE_PATHS.
PATH_TO_TEST_IMAGES_DIR = 'test_images'
TEST_IMAGE_PATHS = [ os.path.join(PATH_TO_TEST_IMAGES_DIR, 'image{}.jpg'.format(i)) for i in range(1, 3) ]
# Size, in inches, of the output images.输出图片打下,单位为英寸
IMAGE_SIZE = (12, 8)
def run_inference_for_single_image(image, graph):
with graph.as_default():
with tf.Session() as sess:
# Get handles to input and output tensors
ops = tf.get_default_graph().get_operations()
all_tensor_names = {output.name for op in ops for output in op.outputs}
tensor_dict = {}
for key in [
'num_detections', 'detection_boxes', 'detection_scores',
'detection_classes', 'detection_masks'
]:
tensor_name = key + ':0'
if tensor_name in all_tensor_names:
tensor_dict[key] = tf.get_default_graph().get_tensor_by_name(
tensor_name)
if 'detection_masks' in tensor_dict:
# The following processing is only for single image
detection_boxes = tf.squeeze(tensor_dict['detection_boxes'], [0])
detection_masks = tf.squeeze(tensor_dict['detection_masks'], [0])
# Reframe is required to translate mask from box coordinates to image coordinates and fit the image size.
real_num_detection = tf.cast(tensor_dict['num_detections'][0], tf.int32)
detection_boxes = tf.slice(detection_boxes, [0, 0], [real_num_detection, -1])
detection_masks = tf.slice(detection_masks, [0, 0, 0], [real_num_detection, -1, -1])
detection_masks_reframed = utils_ops.reframe_box_masks_to_image_masks(
detection_masks, detection_boxes, image.shape[0], image.shape[1])
detection_masks_reframed = tf.cast(
tf.greater(detection_masks_reframed, 0.5), tf.uint8)
# Follow the convention by adding back the batch dimension
tensor_dict['detection_masks'] = tf.expand_dims(
detection_masks_reframed, 0)
image_tensor = tf.get_default_graph().get_tensor_by_name('image_tensor:0')
# Run inference
output_dict = sess.run(tensor_dict,
feed_dict={image_tensor: np.expand_dims(image, 0)})
# all outputs are float32 numpy arrays, so convert types as appropriate
output_dict['num_detections'] = int(output_dict['num_detections'][0])
output_dict['detection_classes'] = output_dict[
'detection_classes'][0].astype(np.uint8)
output_dict['detection_boxes'] = output_dict['detection_boxes'][0]
output_dict['detection_scores'] = output_dict['detection_scores'][0]
if 'detection_masks' in output_dict:
output_dict['detection_masks'] = output_dict['detection_masks'][0]
return output_dict
for image_path in TEST_IMAGE_PATHS:
printf('#start dectec image....:'+PATH_TO_TEST_IMAGES_DIR)
image = Image.open(image_path)
# the array based representation of the image will be used later in order to prepare the
# result image with boxes and labels on it.
printf('load_image_into_numpy_array....:')
image_np = load_image_into_numpy_array(image)
# Expand dimensions since the model expects images to have shape: [1, None, None, 3]
image_np_expanded = np.expand_dims(image_np, axis=0)
# Actual detection.
printf('run_inference_for_single_image....:')
output_dict = run_inference_for_single_image(image_np, detection_graph)
# Visualization of the results of a detection.
printf("visualize_boxes_and_labels_on_image_array")
vis_util.visualize_boxes_and_labels_on_image_array(
image_np,
output_dict['detection_boxes'],
output_dict['detection_classes'],
output_dict['detection_scores'],
category_index,
instance_masks=output_dict.get('detection_masks'),
use_normalized_coordinates=True,
line_thickness=8)
plt.figure(figsize=IMAGE_SIZE)
plt.imshow(image_np)
#pylab.show()
retImageName = image_path.split('.')[0]+'_ret'
printf("save fig "+retImageName)
plt.savefig(retImageName)
#print(image_np)
printf('finish dectec image'+image_path)
浣熊预测demo
#####import 各种依赖包
import numpy as np
import os
import six.moves.urllib as urllib
import sys
import tarfile
import tensorflow as tf
import zipfile
import time
import datetime
import matplotlib
matplotlib.use('Agg')
import pylab
from distutils.version import StrictVersion
from collections import defaultdict
from io import StringIO
from PIL import Image
from matplotlib import pyplot as plt
#定义打印函数,记录执行时间
def printf(info):
print(str(datetime.datetime.now())+':'+info)
# This is needed since the notebook is stored in the object_detection folder.
# 防止你不在object_detection目录,因此将当前目录的父目录添加到系统环境变量
sys.path.append("..")
from object_detection.utils import ops as utils_ops
if StrictVersion(tf.__version__) < StrictVersion('1.9.0'):
raise ImportError('Please upgrade your TensorFlow installation to v1.9.* or later!')
##### Env setup
# This is needed to display the images.
# %matplotlib inline
###Object detection imports
#Here are the imports from the object detection module.
from utils import label_map_util
from utils import visualization_utils as vis_util
###Model preparation
# What model to download.
#获取模型文件ssd_mobilenet_v1_coco_2017_11_17.tar.gz下载地址
MODEL_NAME = 'ssd_mobilenet_v1_coco_2017_11_17'
MODEL_FILE = MODEL_NAME + '.tar.gz'
DOWNLOAD_BASE = 'http://download.tensorflow.org/models/object_detection/'
# Path to frozen detection graph. This is the actual model that is used for the object detection.
PATH_TO_FROZEN_GRAPH = '/home/jiadongfeng/tensorflow/dataset/raccoon_dataset/jdf_train/frozen_inference_graph.pb'
# List of the strings that is used to add correct label for each box.
PATH_TO_LABELS = '/home/jiadongfeng/tensorflow/dataset/raccoon_dataset/jdf_data/raccoon_label_map.pbtxt'
#以下代码是从网址下载模型文件,我这里是在当前object_detection目录使用以下命令下载,并解压到当前目录,因此,将下载代码屏蔽掉了
#wget http://download.tensorflow.org/models/object_detection/ssd_mobilenet_v1_coco_2017_11_17.tar.gz
###Download Model
#opener = urllib.request.URLopener()
#print(DOWNLOAD_BASE + MODEL_FILE)
#opener.retrieve(DOWNLOAD_BASE + MODEL_FILE, MODEL_FILE)
#tar_file = tarfile.open(MODEL_FILE)
#print('tar_file:'+tar_file)
#for file in tar_file.getmembers():
# file_name = os.path.basename(file.name)
# print('file:'+file)
# print('file_name:'+file_name)
# if 'frozen_inference_graph.pb' in file_name:
# tar_file.extract(file, os.getcwd())
###Load a (frozen) Tensorflow model into memory
#加载pb模型文件
printf('start load model file:'+PATH_TO_FROZEN_GRAPH)
detection_graph = tf.Graph()
with detection_graph.as_default():
od_graph_def = tf.GraphDef()
with tf.gfile.GFile(PATH_TO_FROZEN_GRAPH, 'rb') as fid:#object_detection/frozen_inference_graph.pb文件
serialized_graph = fid.read()
od_graph_def.ParseFromString(serialized_graph)
tf.import_graph_def(od_graph_def, name='')
printf('finsh load model file...')
###Loading label map
#PATH_TO_LABELS :#object_detection/data/mscoco_label_map.pbtxt
printf('start load label map file:'+PATH_TO_LABELS)
category_index = label_map_util.create_category_index_from_labelmap(PATH_TO_LABELS, use_display_name=True)
printf('finish load label map file...')
###Helper code
#将图片数据抓换成array数组形式
def load_image_into_numpy_array(image):
(im_width, im_height) = image.size
return np.array(image.getdata()).reshape(
(im_height, im_width, 3)).astype(np.uint8)
printf('finish convert image data to array...')
###Detection
# For the sake of simplicity we will use only 2 images:
# image1.jpg
# image2.jpg
# If you want to test the code with your images, just add path to the images to the TEST_IMAGE_PATHS.
PATH_TO_TEST_IMAGES_DIR = 'test_images'
TEST_IMAGE_PATHS = [ '/home/jiadongfeng/tensorflow/dataset/raccoon_dataset/jdf_images/raccoon-23.jpg' ]
# Size, in inches, of the output images.输出图片打下,单位为英寸
IMAGE_SIZE = (12, 8)
def run_inference_for_single_image(image, graph):
with graph.as_default():
with tf.Session() as sess:
# Get handles to input and output tensors
ops = tf.get_default_graph().get_operations()
all_tensor_names = {output.name for op in ops for output in op.outputs}
tensor_dict = {}
for key in [
'num_detections', 'detection_boxes', 'detection_scores',
'detection_classes', 'detection_masks'
]:
tensor_name = key + ':0'
if tensor_name in all_tensor_names:
tensor_dict[key] = tf.get_default_graph().get_tensor_by_name(
tensor_name)
if 'detection_masks' in tensor_dict:
# The following processing is only for single image
detection_boxes = tf.squeeze(tensor_dict['detection_boxes'], [0])
detection_masks = tf.squeeze(tensor_dict['detection_masks'], [0])
# Reframe is required to translate mask from box coordinates to image coordinates and fit the image size.
real_num_detection = tf.cast(tensor_dict['num_detections'][0], tf.int32)
detection_boxes = tf.slice(detection_boxes, [0, 0], [real_num_detection, -1])
detection_masks = tf.slice(detection_masks, [0, 0, 0], [real_num_detection, -1, -1])
detection_masks_reframed = utils_ops.reframe_box_masks_to_image_masks(
detection_masks, detection_boxes, image.shape[0], image.shape[1])
detection_masks_reframed = tf.cast(
tf.greater(detection_masks_reframed, 0.5), tf.uint8)
# Follow the convention by adding back the batch dimension
tensor_dict['detection_masks'] = tf.expand_dims(
detection_masks_reframed, 0)
image_tensor = tf.get_default_graph().get_tensor_by_name('image_tensor:0')
# Run inference
output_dict = sess.run(tensor_dict,
feed_dict={image_tensor: np.expand_dims(image, 0)})
# all outputs are float32 numpy arrays, so convert types as appropriate
output_dict['num_detections'] = int(output_dict['num_detections'][0])
output_dict['detection_classes'] = output_dict[
'detection_classes'][0].astype(np.uint8)
output_dict['detection_boxes'] = output_dict['detection_boxes'][0]
output_dict['detection_scores'] = output_dict['detection_scores'][0]
if 'detection_masks' in output_dict:
output_dict['detection_masks'] = output_dict['detection_masks'][0]
return output_dict
for image_path in TEST_IMAGE_PATHS:
printf('#start dectec image....:'+PATH_TO_TEST_IMAGES_DIR)
image = Image.open(image_path)
# the array based representation of the image will be used later in order to prepare the
# result image with boxes and labels on it.
printf('load_image_into_numpy_array....:')
image_np = load_image_into_numpy_array(image)
# Expand dimensions since the model expects images to have shape: [1, None, None, 3]
image_np_expanded = np.expand_dims(image_np, axis=0)
# Actual detection.
printf('run_inference_for_single_image....:')
output_dict = run_inference_for_single_image(image_np, detection_graph)
# Visualization of the results of a detection.
printf("visualize_boxes_and_labels_on_image_array")
vis_util.visualize_boxes_and_labels_on_image_array(
image_np,
output_dict['detection_boxes'],
output_dict['detection_classes'],
output_dict['detection_scores'],
category_index,
instance_masks=output_dict.get('detection_masks'),
use_normalized_coordinates=True,
line_thickness=8)
plt.figure(figsize=IMAGE_SIZE)
plt.imshow(image_np)
#pylab.show()
retImageName = image_path.split('.')[0]+'_ret'
printf("save fig "+retImageName)
plt.savefig(retImageName)
#print(image_np)
printf('finish dectec image'+image_path)