# 导入 YOLO 类，用于实例化 YOLO 对象
from ultralytics import YOLO

# 定义 init_model 函数，用于根据传入的模型名称初始化 YOLO 对象
def init_model(model_path):
    model =  YOLO(model_path)
    return model  # 返回 YOLO 对象实例

# 定义 process_frame 函数，用于对输入的视频帧进行目标检测处理
def process_frame(model, frame, show_box, show_mask):

    # 使用传入的 YOLO 模型对输入的视频帧进行目标检测，设置目标检测置信度阈值为 0.25
    # conf = 0.25 使得模型只返回置信度大于 0.25 的检测结果
    # iou = 0.8 使得模型只返回 IoU 大于 0.7 的检测结果，这样可以保留一些重叠较大的检测结果
    results = model.predict(frame, conf=0.25, iou=0.8)

    # 使用检测结果对输入帧进行绘制，按需绘制边框、遮罩及置信度
    processed_frame = results[0].plot(boxes=show_box, masks=show_mask)
    return processed_frame # 返回处理后的帧