目标检测之Tensorflow图像检测

# 载入套件
import os
import pathlib
import tensorflow as tf
import cv2
from os import listdir, path
from os.path import isfile, join
import warnings
import time
from object_detection.utils import label_map_util, config_util
from object_detection.utils import visualization_utils as viz_utils
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
from object_detection.builders import model_builder

# GPU 记忆体配置设定
# GPU 设定为 记忆体动态调整 (dynamic memory allocation)
# 通过 tf.config.experimental.list_physical_devices，我们可以获得当前主机上某种特定运算设备类型（如 GPU 或 CPU ）的列表
gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    # 可以通过 tf.config.experimental.set_memory_growth 将 GPU 的显存使用策略设置为 “仅在需要时申请显存空间”
    tf.config.experimental.set_memory_growth(gpu, True)

#  GPU 设定为固定为 2GB
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:  # 1024*2：2048MB = 2GB
    # 以下代码通过 tf.config.experimental.set_virtual_device_configuration 选项并传入
    # tf.config.experimental.VirtualDeviceConfiguration 实例，设置 TensorFlow 固定消耗 GPU:0 的 1GB 显存（其实可以理解
    # 为建立了一个显存大小为 1GB 的 “虚拟 GPU”）
    tf.config.experimental.set_virtual_device_configuration(gpus[0],
        [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024*2)])


# --------------------------- 方法1：Object Detection From TF2 Saved Model
# 下载模型，并解压缩
def download_model(model_name, model_date):
    base_url = 'http://download.tensorflow.org/models/object_detection/tf2/'
    model_file = model_name + '.tar.gz'
    # 解压缩 从一个 URL 下载文件，如果它不存在缓存中。返回的是下载文件的路径
    model_dir = tf.keras.utils.get_file(fname=model_name,                                 # 文件的名称
                                        origin=base_url + model_date + '/' + model_file,  # 文件的url
                                        untar=True)                                       # 是否需要解压
    return str(model_dir)


MODEL_DATE = '20200711'
MODEL_NAME = 'centernet_hg104_1024x1024_coco17_tpu-32'
# 下载目标检测所使用的模型，并返回存储路径，一般用keras下载下来的文件存储在 C:\Users\用户名\.keras 下
PATH_TO_MODEL_DIR = download_model(MODEL_NAME, MODEL_DATE)
print(PATH_TO_MODEL_DIR)

# 检查模型目录
# 读取 PATH_TO_MODEL_DIR 目录下所有目录及档案
for f in listdir(PATH_TO_MODEL_DIR):
    print(f)

# 从下载的目录载入模型
# 不显示警告讯息
# 通过警告过滤器进行控制是否发出警告消息
warnings.filterwarnings('ignore')   # Suppress warnings

# 从下载的目录载入模型，耗时甚久
PATH_TO_SAVED_MODEL = PATH_TO_MODEL_DIR + "/saved_model"
print('载入模型...', end='')
start_time = time.time()
# 载入模型
detect_fn = tf.saved_model.load(PATH_TO_SAVED_MODEL)
end_time = time.time()
elapsed_time = end_time - start_time
print(f'共花费 {elapsed_time} 秒.')


# 建立 Label 的对照表
# 下载 labels file
def download_labels(filename):
    base_url = 'https://raw.githubusercontent.com/tensorflow/models'
    base_url += '/master/research/object_detection/data/'
    label_dir = tf.keras.utils.get_file(fname=filename,              # 文件的名称
                                        origin=base_url + filename,  # 文件的原始url
                                        untar=False)                 # 文件是否需要解压
    label_dir = pathlib.Path(label_dir)
    return str(label_dir)


LABEL_FILENAME = 'mscoco_label_map.pbtxt'
PATH_TO_LABELS = download_labels(LABEL_FILENAME)  # 下载 Label 对照表文件
print(PATH_TO_LABELS)
# 建立 Label 的对照表 (代码与名称) 标签映射将索引号对应于类别名称，因此假设当我们的卷积网络预测5时，我们知道这对应于飞机。
category_index = label_map_util.create_category_index_from_labelmap(PATH_TO_LABELS, use_display_name=True)

# 选一张图片进行物件侦测
image_np = np.array(Image.open('./data/zebra.jpg'))
# 转为 TensorFlow tensor 资料型态
input_tensor = tf.convert_to_tensor(image_np)
# 加一维，变为 (笔数, 宽, 高, 颜色)
input_tensor = input_tensor[tf.newaxis, ...]

# detections：物件资讯 内含 (候选框, 类别, 机率)
detections = detect_fn(input_tensor)
# 得到检测到的目标数
num_detections = int(detections.pop('num_detections'))
print(f'检测到的物件个数：{num_detections}')
# 转换数据类型
detections = {key: value[0, :num_detections].numpy() for key, value in detections.items()}
detections['num_detections'] = num_detections
# 转为整数
detections['detection_classes'] = detections['detection_classes'].astype(np.int64)

# 打印检测到的结果
print(f'物件资讯 (候选框, 类别, 机率)：')
for detection_boxes, detection_classes, detection_scores in \
        zip(detections['detection_boxes'], detections['detection_classes'], detections['detection_scores']):
    print(np.around(detection_boxes, 4), detection_classes, round(detection_scores*100, 2))

# 筛选Bounding Box，并将图片的物件加框
image_np_with_detections = image_np.copy()
# 加框 Overlay labeled boxes on an image with formatted scores and label name
viz_utils.visualize_boxes_and_labels_on_image_array(
      image_np_with_detections,
      detections['detection_boxes'],
      detections['detection_classes'],
      detections['detection_scores'],
      category_index,                   # a dict containing category dictionaries keyed by category indices
      use_normalized_coordinates=True,  # whether boxes is to be interpreted as normalized coordinates or not.
      max_boxes_to_draw=200,            # maximum number of boxes to visualize.  If None, draw all boxes.
      min_score_thresh=.30,             # minimum score threshold for a box to be visualized
      agnostic_mode=False)              # boolean (default: False) controlling whether to evaluate in class-agnostic mode or not.  This mode will display scores but ignore classes.

# 发现用matplotlib显示无效，但是可以用opencv显示
plt.figure(figsize=(12, 8))
plt.imshow(image_np_with_detections, cmap='viridis')
# 显示处理后的图片
# 存档
saved_file = './data/zebra._detection1.png'
if os.path.isfile(saved_file):
    os.remove(saved_file)
plt.savefig(saved_file)

# -------------------------- 方法2：Object Detection From TF2 Checkpoint
# 从下载的目录载入模型另一种方法，非常快速，快速从下载的目录载入模型
# 设置组态档及模型档路径
PATH_TO_CFG = PATH_TO_MODEL_DIR + "/pipeline.config"
PATH_TO_CKPT = PATH_TO_MODEL_DIR + "/checkpoint"

# 计时开始
print('Loading model... ', end='')
start_time = time.time()
# 载入组态档，再建置模型
# Load pipeline config and build a detection model
configs = config_util.get_configs_from_pipeline_file(PATH_TO_CFG)
model_config = configs['model']
detection_model = model_builder.build(model_config=model_config, is_training=False)
# 还原模型
ckpt = tf.compat.v2.train.Checkpoint(model=detection_model)
# 恢复变量当在其他地方需要为模型重新载入之前保存的参数时，需要再次实例化一个 checkpoint，同时保持键名的一致。再调用 checkpoint 的 restore 方法。
ckpt.restore(os.path.join(PATH_TO_CKPT, 'ckpt-0')).expect_partial()
# 计时完成
end_time = time.time()
elapsed_time = end_time - start_time
print(f'共花费 {elapsed_time} 秒.')


# 任选一张图片进行物件侦测
# 虽然默认的即时执行模式（Eager Execution）为我们带来了灵活及易调试的特性，但在特定的场合，例如追求高性能或部署模型时，我们依然希望
# 使用 TensorFlow 1.X 中默认的图执行模式（Graph Execution），将模型转换为高效的 TensorFlow 图模型。此时，TensorFlow 2 为我们提供
# 了 tf.function 模块，结合 AutoGraph 机制，使得我们仅需加入一个简单的 @tf.function 修饰符，就能轻松将模型以图执行模式运行。
@tf.function
def detect_fn(image):
    image, shapes = detection_model.preprocess(image)
    prediction_dict = detection_model.predict(image, shapes)
    detections = detection_model.postprocess(prediction_dict, shapes)

    return detections


# 读取图档
image_np = np.array(Image.open('./data/zebra.jpg'))
# 转换数据类型
input_tensor = tf.convert_to_tensor(image_np, dtype=tf.float32)
# 增加一维表示样本个数
input_tensor = input_tensor[tf.newaxis, ...]
# 进行检测
detections = detect_fn(input_tensor)
# 获取检测到的目标数
num_detections = int(detections.pop('num_detections'))
print(f'物件个数：{num_detections}')
# 转换数据类型
detections = {key: value[0, :num_detections].numpy() for key, value in detections.items()}
# 显示检测结果
print(f'物件资讯 (候选框, 类别, 机率)：')
for detection_boxes, detection_classes, detection_scores in \
        zip(detections['detection_boxes'], detections['detection_classes'], detections['detection_scores']):
    print(np.around(detection_boxes, 4), int(detection_classes)+1, round(detection_scores*100, 2))

# 结果存入 detections 变数
detections['num_detections'] = num_detections
detections['detection_classes'] = detections['detection_classes'].astype(np.int64)

# 候选框筛选，并将图片的物件加框
# 将物件框起来
# min_score_thresh=.30 表机率(Confidence)至少要大于 30%
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']+1,
      detections['detection_scores'],
      category_index,
      use_normalized_coordinates=True,
      max_boxes_to_draw=200,
      min_score_thresh=.30,
      agnostic_mode=False)

plt.figure(figsize=(12, 8))
plt.imshow(image_np_with_detections, cmap='viridis')

# 显示处理后的图片
# 存档
saved_file = './data/zebra._detection2.png'
if os.path.isfile(saved_file):
    os.remove(saved_file)
plt.savefig(saved_file)