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


# 定义 init_model 函数，用于根据传入的模型名称初始化 YOLO 对象
def init_model(model_path):
    model = YOLO(model_path)  # TODO:  实例化 YOLO 对象，传入模型名称
    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(frame)# TODO

    # 根据输入参数，对指定图片进行目标检测，返回检测结果
    results = model.predict(frame, conf=0.25, iou=0.8)
    processed_frame = results[0].plot(boxes=show_box, masks=show_mask)

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