如何结合Intel OpenVINO预训练模型运行TensorFlow雨伞检测模型？

核心思路 #

先通过预训练人体检测模型定位图像中的所有人，再对每个人体区域用你训练的雨伞模型做二次判断，最终整合两个模型的结果完成标注或提示。

具体实现步骤 #

1. 加载两个OpenVINO模型 #

用OpenVINO Runtime的Core类加载预训练人体检测模型（比如person-detection-retail-0013这类适配场景的模型）和你转换好的雨伞检测模型：

from openvino.runtime import Core
import cv2
import numpy as np

# 初始化OpenVINO核心
ie = Core()

# 加载预训练人体检测模型
person_model = ie.read_model(model="path/to/person-detection-model.xml")
person_compiled_model = ie.compile_model(model=person_model, device_name="CPU")
person_input_layer = person_compiled_model.input(0)
person_output_layer = person_compiled_model.output(0)

# 加载自定义雨伞检测模型
umbrella_model = ie.read_model(model="path/to/your-umbrella-model.xml")
umbrella_compiled_model = ie.compile_model(model=umbrella_model, device_name="CPU")
umbrella_input_layer = umbrella_compiled_model.input(0)
umbrella_output_layer = umbrella_compiled_model.output(0)

2. 预处理输入图像 #

针对两个模型的输入要求分别处理：

def preprocess_person_image(image, input_layer):
    # 按人体模型要求调整尺寸、转换格式
    n, c, h, w = input_layer.shape
    resized_image = cv2.resize(image, (w, h))
    transposed_image = resized_image.transpose(2, 0, 1)  # 转为CHW格式
    input_tensor = np.expand_dims(transposed_image, 0)
    return input_tensor

def preprocess_umbrella_image(image, input_layer):
    # 按你训练雨伞模型时的规则处理，示例为resize+归一化
    n, c, h, w = input_layer.shape
    resized_image = cv2.resize(image, (w, h))
    transposed_image = resized_image.transpose(2, 0, 1)
    normalized_image = (transposed_image / 255.0).astype(np.float32)  # 归一化参数需与训练一致
    input_tensor = np.expand_dims(normalized_image, 0)
    return input_tensor

3. 运行人体检测，获取人体区域 #

解析模型输出，过滤出置信度达标的人体框：

def detect_persons(image, compiled_model, input_layer, output_layer, confidence_threshold=0.5):
    input_tensor = preprocess_person_image(image, input_layer)
    results = compiled_model([input_tensor])[output_layer]
    persons = []
    h, w = image.shape[:2]
    
    # 解析检测结果（格式通常为[image_id, label, confidence, x_min, y_min, x_max, y_max]）
    for result in results[0][0]:
        confidence = result[2]
        if confidence > confidence_threshold:
            # 相对坐标转绝对像素坐标
            x_min = int(result[3] * w)
            y_min = int(result[4] * h)
            x_max = int(result[5] * w)
            y_max = int(result[6] * h)
            persons.append((x_min, y_min, x_max, y_max))
    return persons

4. 对人体区域做雨伞检测 #

遍历每个人体框，裁剪区域后用雨伞模型判断是否持伞：

def check_umbrella_in_person(person_box, image, compiled_model, input_layer, output_layer, umbrella_threshold=0.7):
    x_min, y_min, x_max, y_max = person_box
    # 裁剪人体区域
    person_region = image[y_min:y_max, x_min:x_max]
    # 预处理后推理
    input_tensor = preprocess_umbrella_image(person_region, input_layer)
    result = compiled_model([input_tensor])[output_layer]
    # 假设模型为二分类，索引1对应持伞置信度（需根据你的训练输出调整）
    umbrella_confidence = result[0][1]
    return umbrella_confidence > umbrella_threshold

5. 整合结果并可视化 #

结合两个模型的输出，完成标注或提示：

def process_image(image_path):
    image = cv2.imread(image_path)
    if image is None:
        print("无法读取图像")
        return
    
    # 检测图像中的人体
    persons = detect_persons(image, person_compiled_model, person_input_layer, person_output_layer)
    has_umbrella_person = False
    
    # 逐个检查人体是否持伞
    for person_box in persons:
        is_holding_umbrella = check_umbrella_in_person(person_box, image, umbrella_compiled_model, umbrella_input_layer, umbrella_output_layer)
        if is_holding_umbrella:
            has_umbrella_person = True
            # 绘制标注框
            x_min, y_min, x_max, y_max = person_box
            cv2.rectangle(image, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
            cv2.putText(image, "Umbrella", (x_min, y_min-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,255,0), 2)
    
    # 输出结果
    if has_umbrella_person:
        cv2.imshow("Result", image)
        cv2.waitKey(0)
    else:
        print("无人持伞")
    
    cv2.destroyAllWindows()

# 运行示例
process_image("test_image.jpg")

注意事项 #

• 预训练人体模型选择：优先选适配你场景的模型，比如person-detection-retail-0013适合近距离场景，person-detection-0200适合通用场景，确保输出为标准检测框格式
• 雨伞模型输出解析：根据你训练时的设置调整，若为目标检测模型，需判断人体区域内是否存在伞的检测框；若为二分类，需对应正确的类别索引
• 阈值调整：根据实际测试效果，微调人体检测和雨伞检测的置信度阈值，平衡误检与漏检

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