AI 社交距离检测器：使用 TensorFlow 和 MobileNet 进行对象检测
3.32/5 (7投票s)
在本文中，我们将使用绘图函数来描绘检测到的对象。
在本文中，我们将继续学习如何使用人工智能构建社交距离检测器。具体来说，我们将学习如何使用 TensorFlow 检测图像中的对象。然后，我们将使用之前开发的方法标注检测到的对象。
阅读完本文后，您将获得生成下图结果的 Python 代码。配套代码在此处：这里。
必备组件
我首先从这里下载了一个预训练的 TensorFlow 模型。该软件包包含两个文件：
- detect.tflite – 以 TensorFlow Lite 格式保存的预训练模型。我们将使用此文件执行对象检测。
- labelmap.txt – 一个包含检测到对象标签的文本文件。
所选模型 (detect.tflite) 返回有关检测到的对象的几条信息，包括：
- 边界框 – 围绕检测到的对象的矩形。
- 类别 ID – 描述检测到的对象的类别标识符。给定此 ID，您可以从 labelmap.txt 文件中获取对象标签。
- 置信度 – 检测置信度，通常表示为 0 到 1 之间的浮点数。值越大，置信度越高。
入门
鉴于以上信息，我首先实现了 Inference 类 (inference.py)，在该类中我定义了两个方法：load_model_and_configure 和 load_labels_from_file。这两个方法都在 Inference 类的初始化器中调用。第一个方法 load_model_and_configure 如下所示：
from tensorflow import lite as tflite
def load_model_and_configure(self, model_path):
""" Loads the model and configures input image dimensions accordingly
: model_path: A full path to the file containing the model
"""
# Load model from file
self.interpreter = tflite.Interpreter(model_path)
# Allocate tensors
self.interpreter.allocate_tensors()
# Get input and output details
self.input_details = self.interpreter.get_input_details()
self.output_details = self.interpreter.get_output_details()
# Store input image dimensions
self.input_image_height = self.input_details[0]['shape'][1]
self.input_image_width = self.input_details[0]['shape'][2]
load 方法接受一个参数 model_path。它指向 *.tflite 文件的位置。然后，我使用 TensorFlow Lite 中的 Interpreter 类加载此文件。我将生成的对象存储在 Inference 类的 interpreter 字段中。然后，我分配内存 (allocate_tensors)，并获取模型的输入和输出。为此，我使用了 Interpreter 类中的 get_input_details 和 get_output_details。最后，我使用输入详细信息读取模型输入可接受图像的尺寸。模型可能经过预训练，只能接受特定格式的图像。根据模型，您可以使用解释器类的 resize_tensor_input 来更改输入大小。在此，我使用 prepare_image 方法（见下文）将输入图像调整为模型要求的大小。该方法将颜色通道的排列从 BGR（蓝-绿-红）转换为 RGB，然后将图像调整为模型期望的尺寸。这两个步骤都使用 OpenCV 完成。
def prepare_image(self, image):
""" Prepares image for the TensorFlow inference
: image: An input image
"""
# Convert image to BGR
image = opencv.cvtColor(image, opencv.COLOR_BGR2RGB)
# Get new size
new_size = (self.input_image_height, self.input_image_width)
# Resize
image = opencv.resize(image, new_size, interpolation = opencv.INTER_AREA)
return image
为了从文件中读取对象标签，我实现了以下方法：
def load_labels_from_file(self, file_path):
with open(file_path, 'r') as file:
self.labels = [line.strip() for line in file.readlines()]
我读取文件内容并将生成的字符串集合存储在 Inference 类的 labels 字段中。
我使用上述方法和 load_model_and_configure 来初始化 Inference 类，如下所示。这确保模型和标签在执行对象检测之前已准备就绪。
class Inference(object):
def __init__(self, model_file_path, labels_file_path):
# Load model
self.load_model_and_configure(model_file_path)
# Load labels
self.load_labels_from_file(labels_file_path)
目标检测
模型和标签准备就绪后，我们就可以执行对象检测了。这里的通用流程是，您首先将输入图像或图像序列设置为输入张量。您可以通过 Interpreter 类实例的 tensor 成员来实现这一点。这是 Inference 类中的示例：
def set_input_tensor(self, image):
tensor_index = self.input_details[0]['index']
input_tensor = self.interpreter.tensor(tensor_index)()[0]
input_tensor[:,:] = image
此方法读取输入张量索引，然后将张量设置为传递给参数的图像。
下一步是使用 Interpreter 类的 invoke 方法执行推理，并从输出张量读取推理结果。输出详细信息的数量因模型而异。我在这里使用的模型返回边界框、类别标识符和置信度。我使用 get_output_details 辅助方法读取它们（它返回给定索引处的张量）。
def get_output_tensor(self, index):
tensor = self.interpreter.get_tensor(self.output_details[index]['index'])
return np.squeeze(tensor)
为了处理检测对象的完整工作流程，我实现了以下方法：
def detect_objects(self, image, threshold):
"""Returns a list of detection results"""
# Store input image size
input_image_size = image.shape[-2::-1]
# Prepare image
image = self.prepare_image(image)
# Set image as the input tensor
self.set_input_tensor(image)
# Perform inference
self.interpreter.invoke()
# Get all output details
boxes = self.get_output_tensor(0)
classes = self.get_output_tensor(1)
scores = self.get_output_tensor(2)
# Filter out detections below the threshold
results = []
for i in range(scores.size):
if scores[i] >= threshold:
result = {
'rectangle': self.convert_bounding_box_to_rectangle_points(
boxes[i], input_image_size),
'label': self.labels[int(classes[i])],
}
results.append(result)
# Return informations about detected objects
return results
在获得推理结果集合后，我们通过删除所有得分低于阈值的检测来过滤检测结果。此阈值来自 detect_objects 方法的最后一个参数。因此，detect_objects 方法将返回一个对象列表，每个对象都有两个属性：rectangle 和 label。
另请注意，TensorFlow 模型返回的边界框（四元素数组）以归一化单位表示。左上角（数组的前两个元素）和右下角（第三个和第四个元素）的坐标具有分数。它们除以输入图像（传递给模型的图像）的高度和宽度。因此，要恢复边界框在输入图像中的位置，我们需要将值乘以输入图像的宽度和高度。我使用以下方法完成此操作：
def convert_bounding_box_to_rectangle_points(self, bounding_box, input_image_size):
width = input_image_size[0]
height = input_image_size[1]
top_left_corner = (int(bounding_box[1] * width), int(bounding_box[0] * height))
bottom_right_corner = (int(bounding_box[3] * width), int(bounding_box[2] * height))
return (top_left_corner, bottom_right_corner)
该方法返回两个值：top_left_corner 和 bottom_right_corner。然后，我使用它们通过 OpenCV 绘制矩形，如下所述。
显示结果
运行推理后，我标注了检测到的对象。为此，我使用了 OpenCV 的 rectangle 和 putText 方法，这些方法在上一篇文章中已作介绍。相应的函数在 ImageHelper 类（请参阅配套代码）的一个静态方法 draw_rectangle_and_label 中实现，该方法接收单个推理结果并在图像上绘制矩形和标签。我为 detect_objects 返回的列表中的每个元素调用此方法。最后，如下所示，我使用 OpenCV 的 imshow 函数显示图像。
def display_image_with_detected_objects(image, inference_results):
# Prepare window
opencv.namedWindow(common.WINDOW_NAME, opencv.WINDOW_GUI_NORMAL)
# Draw rectangles and labels on the image
for i in range(len(inference_results)):
current_result = inference_results[i]
ImageHelper.draw_rectangle_and_label(
image, current_result['rectangle'], current_result['label'])
# Display image
opencv.imshow(common.WINDOW_NAME, image)
# Wait until user presses any key
opencv.waitKey(0)
整合
现在，我们可以将以上所有部分组合起来，编写应用程序的主部分（请参阅配套代码中的 main.py）。
import common
from image_helper import ImageHelper as imgHelper
from inference import Inference as model
if __name__ == "__main__":
# Load and prepare model
model_file_path = '../Models/01_model.tflite'
labels_file_path = '../Models/02_labels.txt'
# Initialize model
ai_model = model(model_file_path, labels_file_path)
# Get input image
image = imgHelper.load_image('../Images/Lena.png')
# Detect objects
score_threshold = 0.5
results = ai_model.detect_objects(image, score_threshold)
# Display results
imgHelper.display_image_with_detected_objects(image, results)
如上所示，我们从 Models 子文件夹加载模型和标签。然后，我们配置模型，加载图像，并以 0.5 的阈值执行推理。这意味着置信度低于 50% 的检测将被拒绝。最后，我们显示带有已标注检测对象的图像。运行 main.py 后，您将获得前面描绘的结果。
摘要
我们学习了如何使用 Python 中的预训练 MobileNet 模型执行对象检测。我们实现了从加载模型到解释和显示结果的完整流程。现在，我们可以检测图像中的更多人物并计算他们之间的距离。我们将在下一篇文章中介绍这一点。