181 lines
7.8 KiB
Python
181 lines
7.8 KiB
Python
# -*- coding: utf-8 -*-
|
||
|
||
import cv2
|
||
import numpy as np
|
||
from ultralytics import YOLO
|
||
import random
|
||
import os
|
||
from tqdm import tqdm # 引入进度条库
|
||
|
||
# --- 配置参数 ---
|
||
|
||
# 模型文件路径
|
||
# MODEL_PATH = os.path.join('YOLOv9e-seg', 'weights', 'best.pt')
|
||
MODEL_PATH = os.path.join('YOLO11l-seg', 'best.pt')
|
||
# 输入视频文件路径
|
||
VIDEO_IN_PATH = 'LC_Video_1.mp4'
|
||
# 输出视频文件路径
|
||
VIDEO_OUT_PATH = 'output_a_segmented_fps_controlled.mp4'
|
||
|
||
# --- 新增配置 ---
|
||
# 要处理的视频帧率。例如，设置为 5 表示每秒大约只对 5 帧进行模型推理。
|
||
# 设置为 None 则处理所有帧。
|
||
PROCESS_FPS = None
|
||
|
||
# 推理参数
|
||
CONFIDENCE_THRESHOLD = 0.3 # 置信度阈值，低于此值的检测将被忽略
|
||
MASK_ALPHA = 0.4 # 分割掩码的透明度 (0.0 完全透明, 1.0 完全不透明)
|
||
|
||
# --- 主处理逻辑 ---
|
||
|
||
def main():
|
||
"""
|
||
主函数，执行视频分割的完整流程。
|
||
"""
|
||
# 1. 加载预训练的 YOLOv9 分割模型
|
||
# YOLO() 类会自动处理模型的加载、权重初始化以及设备选择（优先使用可用的 GPU）
|
||
print(f"正在加载模型: {MODEL_PATH}...")
|
||
try:
|
||
model = YOLO(MODEL_PATH)
|
||
except Exception as e:
|
||
print(f"错误: 无法加载模型。请检查路径是否正确以及 Ultralytics 是否已正确安装。")
|
||
print(f"详细错误: {e}")
|
||
return
|
||
|
||
# 获取模型能够识别的所有类别名称
|
||
class_names = model.names
|
||
print(f"模型已加载，可识别 {len(class_names)} 个类别。")
|
||
|
||
# 2. 初始化视频读写对象
|
||
# 使用 OpenCV 打开输入视频文件
|
||
cap = cv2.VideoCapture(VIDEO_IN_PATH)
|
||
if not cap.isOpened():
|
||
print(f"错误: 无法打开视频文件: {VIDEO_IN_PATH}")
|
||
return
|
||
|
||
# 获取输入视频的属性（帧率、宽度、高度、总帧数）
|
||
fps = cap.get(cv2.CAP_PROP_FPS)
|
||
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) # 获取总帧数用于进度条
|
||
print(f"输入视频属性: {frame_width}x{frame_height} @ {fps:.2f} FPS, 总共 {total_frames} 帧")
|
||
|
||
# 定义视频编码器并创建 VideoWriter 对象以保存输出视频
|
||
# 'mp4v' 是适用于.mp4 文件的常用编码器
|
||
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
|
||
out = cv2.VideoWriter(VIDEO_OUT_PATH, fourcc, fps, (frame_width, frame_height))
|
||
|
||
# 3. 为每个类别生成一个稳定的随机颜色
|
||
# 这确保了在整个视频中，同一类别的对象掩码颜色保持一致
|
||
random.seed(42) # 固定随机种子以保证每次运行颜色相同
|
||
class_colors = {
|
||
name: [random.randint(0, 255) for _ in range(3)]
|
||
for name in class_names.values()
|
||
}
|
||
|
||
# --- 新增：计算帧处理间隔 ---
|
||
frame_counter = 0
|
||
frame_skip_interval = 1
|
||
if PROCESS_FPS is not None and PROCESS_FPS > 0:
|
||
frame_skip_interval = round(fps / PROCESS_FPS)
|
||
# 确保至少为1，避免除零错误
|
||
if frame_skip_interval < 1:
|
||
frame_skip_interval = 1
|
||
print(f"处理帧率设置为 {PROCESS_FPS} FPS. 将每 {frame_skip_interval} 帧运行一次模型推理。")
|
||
else:
|
||
print("将处理所有帧。")
|
||
|
||
# 4. 逐帧处理视频，并使用 tqdm 显示进度条
|
||
print("开始逐帧处理视频... (在预览窗口激活时按 'q' 键可随时退出)")
|
||
for _ in tqdm(range(total_frames), desc="处理进度"):
|
||
# 从视频中读取一帧
|
||
success, frame = cap.read()
|
||
|
||
if not success:
|
||
# 如果视频读取结束，则跳出循环
|
||
print("\n视频读取完毕或发生错误。")
|
||
break
|
||
|
||
frame_counter += 1
|
||
|
||
# 5. 对当前帧执行 YOLOv9 推理 (有条件地)
|
||
# 仅在达到处理间隔时才运行模型，否则将重用上一次的结果
|
||
if frame_counter % frame_skip_interval == 0:
|
||
results = model.predict(source=frame, conf=CONFIDENCE_THRESHOLD, verbose=False)
|
||
|
||
# 检查是否有任何有效的 'results' 可供绘制
|
||
# 'results' 可能来自当前帧的推理，或来自之前的帧
|
||
# `locals()` 检查 `results` 变量是否已在当前作用域中定义
|
||
if 'results' in locals() and results is not None:
|
||
# ------------------- FIX START -------------------
|
||
# `results` 是一个列表，因为我们处理的是单张图片，所以需要取出第一个元素
|
||
result = results[0]
|
||
# -------------------- FIX END --------------------
|
||
|
||
# 6. 可视化：绘制半透明掩码和边界框
|
||
|
||
# 创建一个与原始帧相同的副本，用于绘制掩码。这是实现透明效果的关键步骤。
|
||
overlay = frame.copy()
|
||
|
||
# 检查是否存在检测到的掩码
|
||
if result.masks is not None:
|
||
# 遍历每个检测到的对象
|
||
# zip() 函数将边界框和掩码一一对应起来
|
||
for box, mask in zip(result.boxes, result.masks):
|
||
# 获取类别 ID 和类别名称
|
||
class_id = int(box.cls)
|
||
class_name = class_names[class_id]
|
||
|
||
# 获取该类别的预定义颜色
|
||
color = class_colors[class_name]
|
||
|
||
# --- 绘制分割掩码 ---
|
||
# [修正] mask.xy返回的是一个列表的列表，需要用np.array处理
|
||
points = np.array(mask.xy, dtype=np.int32)
|
||
# 在 overlay 层上填充多边形（完全不透明）
|
||
cv2.fillPoly(overlay, [points], color)
|
||
|
||
# --- 绘制边界框和标签 ---
|
||
# 获取边界框坐标
|
||
x1, y1, x2, y2 = map(int, box.xyxy[0]) # box.xyxy也是列表，取第一个
|
||
# 在原始帧上绘制矩形边界框
|
||
cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
|
||
|
||
# 准备标签文本（类别名 + 置信度）
|
||
confidence = float(box.conf)
|
||
label = f"{class_name}: {confidence:.2f}"
|
||
|
||
# 计算文本尺寸以绘制背景
|
||
(text_width, text_height), baseline = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
|
||
# 绘制文本背景框
|
||
cv2.rectangle(frame, (x1, y1 - text_height - baseline), (x1 + text_width, y1), color, -1)
|
||
# 在背景框上放置文本
|
||
cv2.putText(frame, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
|
||
|
||
# 7. 融合原始帧和掩码层
|
||
# cv2.addWeighted() 将 overlay 和 frame 按照指定的权重（透明度）混合
|
||
# 公式为: output = frame * (1 - alpha) + overlay * alpha + 0
|
||
cv2.addWeighted(overlay, MASK_ALPHA, frame, 1 - MASK_ALPHA, 0, frame)
|
||
|
||
# 8. 将处理后的帧写入输出视频文件
|
||
out.write(frame)
|
||
|
||
# # 9. [新增] 实时显示处理后的帧
|
||
# cv2.imshow('Real-time Segmentation Preview', frame)
|
||
|
||
# # 检测 'q' 键是否被按下，如果是则退出循环
|
||
# # cv2.waitKey(1) 对于视频处理至关重要，它等待 1ms，允许 OpenCV 刷新窗口
|
||
# if cv2.waitKey(1) & 0xFF == ord('q'):
|
||
# print("\n用户中断处理。")
|
||
# break
|
||
|
||
# 10. 释放资源
|
||
cap.release()
|
||
out.release()
|
||
# cv2.destroyAllWindows()
|
||
print(f"处理完成！")
|
||
print(f"输出视频已保存至: {VIDEO_OUT_PATH}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |