import cv2
from ultralytics import YOLO

# ------------------- 配置参数 -------------------
VIDEO_PATH = '1234520251015_152225.mp4'
WEIGHTS_PATH = 'yolov8n-seg.pt'
OUTPUT_PATH = 'output_segmented_video.mp4'  # 输出视频路径
CONFIDENCE_THRESHOLD = 0.5  # 置信度阈值
SHOW_VIDEO = True  # 是否实时显示（可设为 False 以加快处理）
# ------------------------------------------------

# 加载 YOLOv8 分割模型
model = YOLO(WEIGHTS_PATH)  # 自动加载模型结构和权重

# 打开视频文件
cap = cv2.VideoCapture(VIDEO_PATH)

# 获取视频的宽高、帧率
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = int(cap.get(cv2.CAP_PROP_FPS))

# 定义视频写入器
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter(OUTPUT_PATH, fourcc, fps, (width, height))

# 检查视频是否成功打开
if not cap.isOpened():
    print("❌ 无法打开视频文件！请检查路径是否正确。")
    exit()

print("✅ 开始处理视频...")

while True:
    ret, frame = cap.read()
    if not ret:
        print("🔚 视频读取完成或结束。")
        break

    # 使用 YOLOv8 进行推理（分割）
    results = model(frame, conf=CONFIDENCE_THRESHOLD, imgsz=640)

    # 在原图上绘制分割结果
    annotated_frame = results[0].plot()  # plot() 方法会自动画出边界框、分割掩码、标签等

    # 写入输出视频
    out.write(annotated_frame)

    # 实时显示（可选）
    if SHOW_VIDEO:
        cv2.imshow('YOLOv8 Segmentation', annotated_frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):  # 按 'q' 键退出
            print("🛑 用户中断。")
            break

# 释放资源
cap.release()
out.release()
cv2.destroyAllWindows()

print(f"✅ 处理完成！分割后的视频已保存为：{OUTPUT_PATH}")