TensorFlow2.x目标检测API测试代码使用演示

Color Space

发布于 2020-11-09 14:26:44

1.9K0

发布于 2020-11-09 14:26:44

文章被收录于专栏：OpenCV与AI深度学习OpenCV与AI深度学习

TensorFlow2.x Object Detection API 的安装与配置可参考前面的两篇文章：

TensorFlow2.x GPU版安装与CUDA版本选择指南

TensorFlow2.x 目标检测API安装配置步骤详细教程

安装配置完成后，可以使用代码测试了。

一、在Model Zoo下载需要测试的模型，这里选择的SSD MobileNet V2 FPNLite 320x320

https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/tf2_detection_zoo.md

下载后解压，可以看到有如些这些文件(这里后放到D:\TensorFlow\Test\model文件夹下)

二、在Object Detection API安装目录找到pbtxt配置文件，D:\TensorFlow\models\research\object_detection\data

将mscoco_label_map.pbtxt拷贝到指定文件夹，这里放到model文件夹内与saved_model文件夹同目录

三、使用测试图像，加载模型测试，如果缺cv2模块则pip install opencv-python，单张图片测试代码如下：

#!/usr/bin/env python
# coding: utf-8
"""
Object Detection From TF2 Saved Model
=====================================
"""
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'    # Suppress TensorFlow logging (1)
import pathlib
import tensorflow as tf
import time
import cv2
import numpy as np
from object_detection.utils import label_map_util
from object_detection.utils import visualization_utils as viz_utils

tf.get_logger().setLevel('ERROR')           # Suppress TensorFlow logging (2)

# Enable GPU dynamic memory allocation
gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

model_dir = "./model"

label_path = './model/mscoco_label_map.pbtxt'

path_saved_model = model_dir + "/saved_model"

# Load saved model and build the detection function
detect_fn = tf.saved_model.load(path_saved_model)
category_index = label_map_util.create_category_index_from_labelmap(label_path,use_display_name=True)

import warnings
warnings.filterwarnings('ignore')   # Suppress Matplotlib warnings


#----------------read image and test--------------------#
image_path = "./images/6.jpg"
image_np = cv2.imread(image_path)
input_tensor = tf.convert_to_tensor(image_np)   
input_tensor = input_tensor[tf.newaxis, ...]
detections = detect_fn(input_tensor)

num_detections = int(detections.pop('num_detections'))
detections = {key: value[0, :num_detections].numpy()
               for key, value in detections.items()}
detections['num_detections'] = num_detections

# detection_classes should be ints.
detections['detection_classes'] = detections['detection_classes'].astype(np.int64)

image_np_with_detections = image_np.copy()

viz_utils.visualize_boxes_and_labels_on_image_array(
        image_np_with_detections,
        detections['detection_boxes'],
        detections['detection_classes'],
        detections['detection_scores'],
        category_index,
        use_normalized_coordinates=True,
        max_boxes_to_draw=200,
        min_score_thresh=0.30,
        agnostic_mode=False)

cv2.imshow("object_detection_demo", image_np_with_detections)
cv2.waitKey()
cv2.destroyAllWindows()
print('Done')

测试图像：