基于TensorFlow的口罩检测项目：如何添加语音提示与抓拍功能？

准备工作 #

• 安装音频依赖库：执行 pip install pygame
• 准备提示音频文件（如alert.wav），放在代码同目录下，或填写完整路径
• 确认你的label_map.pbtxt中，未戴口罩对应的类别ID（比如常见配置中是2，需根据实际文件调整）

修改后的完整代码 #

import cv2 
import numpy as np
import pygame
import time
import tensorflow as tf
from object_detection.utils import label_map_util
from object_detection.utils import visualization_utils as viz_utils

# 替换为你的标注文件路径
ANNOTATION_PATH = "path/to/annotations"
category_index = label_map_util.create_category_index_from_labelmap(ANNOTATION_PATH+'/label_map.pbtxt')

# 初始化音频模块与提示音
pygame.mixer.init()
alert_sound = pygame.mixer.Sound("alert.wav")  # 替换为你的音频文件路径

# 冷却时间设置（避免频繁触发）
ALERT_COOLDOWN = 2  # 音频提示间隔（秒）
CAPTURE_COOLDOWN = 5  # 抓拍间隔（秒）
last_alert_time = 0
last_capture_time = 0

# Setup capture
cap = cv2.VideoCapture(0)
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

# 加载检测模型（需确保detect_fn已正确定义）
# detect_fn = tf.saved_model.load(PATH_TO_SAVED_MODEL)

while True: 
    ret, frame = cap.read()
    if not ret:
        break
    image_np = np.array(frame)
    
    input_tensor = tf.convert_to_tensor(np.expand_dims(image_np, 0), dtype=tf.float32)
    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)

    label_id_offset = 1
    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']+label_id_offset,
                detections['detection_scores'],
                category_index,
                use_normalized_coordinates=True,
                max_boxes_to_draw=5,
                min_score_thresh=.7,
                agnostic_mode=False)

    # 检测未戴口罩目标并触发功能
    current_time = time.time()
    for i in range(num_detections):
        score = detections['detection_scores'][i]
        class_id = detections['detection_classes'][i] + label_id_offset
        # 匹配未戴口罩类别（此处ID为2，需根据你的label_map调整）
        if score >= 0.7 and class_id == 2:
            # 音频提示（带冷却）
            if current_time - last_alert_time >= ALERT_COOLDOWN:
                alert_sound.play()
                last_alert_time = current_time
            # 抓拍保存（带冷却）
            if current_time - last_capture_time >= CAPTURE_COOLDOWN:
                capture_filename = f"no_mask_capture_{int(current_time)}.jpg"
                cv2.imwrite(capture_filename, frame)
                print(f"已保存抓拍：{capture_filename}")
                last_capture_time = current_time

    cv2.imshow('Face Mask Detection',  cv2.resize(image_np_with_detections, (960, 720)))

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()
pygame.mixer.quit()  # 退出音频模块

关键说明 #

1. 类别ID调整：务必根据你的label_map.pbtxt文件修改class_id == 2中的数字，比如如果文件中未戴口罩的条目是id: 1，则改为class_id == 1
2. 冷却机制：设置ALERT_COOLDOWN和CAPTURE_COOLDOWN避免短时间内重复触发提示和抓拍，可根据需求调整时长
3. 音频文件：支持wav格式，若用其他格式可参考pygame文档转换或调整加载方式
4. 抓拍路径：默认保存在代码同目录，如需指定路径，修改capture_filename为完整路径即可
5. 模型加载：需确保detect_fn已通过tf.saved_model.load()正确加载你的检测模型

内容的提问来源于stack exchange，提问作者user19339325