491 lines
18 KiB
Python
491 lines
18 KiB
Python
import os
|
||
from pathlib import Path
|
||
import numpy as np
|
||
import supervision as sv
|
||
import cv2
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
|
||
def generate_video_with_boxes(
|
||
boxes_data: list[dict],
|
||
input_video_path: str,
|
||
output_video_path: str,
|
||
frame_annotation_interval: int
|
||
) -> None:
|
||
"""
|
||
将提供的标注数据渲染到视频上，支持按帧间隔进行标注。
|
||
|
||
:param boxes_data: 包含标注信息的列表，每个元素结构为:
|
||
{
|
||
"frame_id": int,
|
||
"boxes": List[Tuple[int, int, int, int, str]] # (x1, y1, x2, y2, label)
|
||
}
|
||
:param input_video_path: 输入视频文件路径
|
||
:param output_video_path: 输出视频文件路径
|
||
:param frame_annotation_interval: 标注间隔（单位：帧），默认为 1（每帧标注）
|
||
:return: 无
|
||
"""
|
||
# -------------------------------------------------
|
||
# 1. 视频读取与基本配置
|
||
# -------------------------------------------------
|
||
cap = cv2.VideoCapture(input_video_path)
|
||
if not cap.isOpened():
|
||
raise FileNotFoundError(f"无法打开视频文件或流: {input_video_path}")
|
||
fps = cap.get(cv2.CAP_PROP_FPS)
|
||
video_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||
video_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||
|
||
# -------------------------------------------------
|
||
# 2. 初始化 VideoWriter
|
||
# -------------------------------------------------
|
||
Path(os.path.dirname(output_video_path)).mkdir(parents=True, exist_ok=True)
|
||
fourcc = cv2.VideoWriter_fourcc(*'mp4v') # type: ignore
|
||
out = cv2.VideoWriter(output_video_path, fourcc, fps, (video_width, video_height))
|
||
|
||
# -------------------------------------------------
|
||
# 3. 初始化标注工具
|
||
# -------------------------------------------------
|
||
box_annotator = sv.BoxAnnotator()
|
||
label_annotator = sv.RichLabelAnnotator(
|
||
font_path="C:/Windows/Fonts/simhei.ttf",
|
||
text_color=sv.Color.WHITE,
|
||
text_padding=5,
|
||
font_size=20
|
||
)
|
||
|
||
# -------------------------------------------------
|
||
# 4. 数据预处理：构建 frame_id -> boxes 的映射
|
||
# -------------------------------------------------
|
||
frame_to_boxes: dict[int, list[tuple[int, int, int, int, str]]] = {}
|
||
for entry in boxes_data:
|
||
frame_id = entry["frame_id"]
|
||
boxes = entry["boxes"]
|
||
# 确保每个框都有正确的结构
|
||
cleaned_boxes = []
|
||
for box in boxes:
|
||
if len(box) == 5:
|
||
cleaned_boxes.append(box) # (x1, y1, x2, y2, label)
|
||
frame_to_boxes[frame_id] = cleaned_boxes
|
||
|
||
# -------------------------------------------------
|
||
# 5. 主循环：逐帧读取并渲染
|
||
# -------------------------------------------------
|
||
frame_idx = 0
|
||
while cap.isOpened():
|
||
ret, frame = cap.read()
|
||
if not ret:
|
||
break
|
||
|
||
# 计算当前帧对应的标注帧索引（考虑间隔）
|
||
annotation_frame_idx = frame_idx // frame_annotation_interval
|
||
|
||
# 如果当前帧没有标注数据，直接写入原始帧
|
||
if annotation_frame_idx not in frame_to_boxes:
|
||
out.write(frame)
|
||
frame_idx += 1
|
||
continue
|
||
|
||
# 获取当前帧的所有框信息
|
||
boxes_info = frame_to_boxes[annotation_frame_idx]
|
||
|
||
boxes = []
|
||
labels = []
|
||
for _, (x1, y1, x2, y2, label) in enumerate(boxes_info):
|
||
# 坐标
|
||
boxes.append([x1, y1, x2, y2])
|
||
# 使用索引作为唯一标识符，或直接使用 label
|
||
labels.append(label)
|
||
|
||
# 转换为 NumPy 数组
|
||
boxes_np = np.array(boxes, dtype=np.float64)
|
||
|
||
# 构建 Detections 对象
|
||
detections = sv.Detections(
|
||
xyxy=boxes_np,
|
||
confidence=np.ones(len(boxes_np)), # 默认置信度为 1.0
|
||
class_id=np.zeros(len(boxes_np), dtype=int) # 类别 ID 在此场景下不重要
|
||
)
|
||
detections.tracker_id = np.arange(len(boxes_np), dtype=int) # 使用索引作为 ID
|
||
|
||
# 绘制边框和标签
|
||
annotated_frame = box_annotator.annotate(scene=frame.copy(), detections=detections)
|
||
annotated_frame = label_annotator.annotate(
|
||
scene=annotated_frame,
|
||
detections=detections,
|
||
labels=labels
|
||
)
|
||
|
||
# 写入帧
|
||
out.write(annotated_frame)
|
||
frame_idx += 1
|
||
|
||
# -------------------------------------------------
|
||
# 6. 资源释放
|
||
# -------------------------------------------------
|
||
cap.release()
|
||
out.release()
|
||
cv2.destroyAllWindows()
|
||
print(f"视频渲染完成，已保存至: {output_video_path}")
|
||
|
||
def process_track_id(
|
||
track_id: int,
|
||
frame_list: list[tuple[int, list[int]]],
|
||
input_video_path: str,
|
||
output_video_root: str,
|
||
frame_width: int,
|
||
frame_height: int,
|
||
target_fps: int,
|
||
frame_interval: int
|
||
) -> str:
|
||
"""
|
||
处理单个track_id，生成对应的视频
|
||
"""
|
||
# 计算需要生成的总帧数（确保覆盖所有物体帧且不少于两秒）
|
||
min_frames_for_2s = 25 * 2 # 2秒 @ 25fps
|
||
object_based_frames = len(frame_list) * frame_interval
|
||
max_output_frame = max(object_based_frames, min_frames_for_2s)
|
||
|
||
# 创建输出视频路径
|
||
output_path = os.path.join(output_video_root, f"{track_id}.mp4")
|
||
|
||
# 尝试使用GPU硬件编码
|
||
try:
|
||
# 对于不同平台的GPU编码，使用不同的fourcc
|
||
# Windows平台使用h264_nvenc或h264_amf
|
||
# 如果GPU编码不可用，会回退到CPU编码
|
||
fourcc = cv2.VideoWriter_fourcc(*'h264') # type: ignore
|
||
out = cv2.VideoWriter(output_path, fourcc, target_fps, (frame_width, frame_height))
|
||
# 检查是否成功打开
|
||
if not out.isOpened():
|
||
# 尝试其他编码方式
|
||
fourcc = cv2.VideoWriter_fourcc(*'mp4v') # type: ignore
|
||
out = cv2.VideoWriter(output_path, fourcc, target_fps, (frame_width, frame_height))
|
||
except Exception:
|
||
# 异常时使用CPU编码
|
||
fourcc = cv2.VideoWriter_fourcc(*'mp4v') # type: ignore
|
||
out = cv2.VideoWriter(output_path, fourcc, target_fps, (frame_width, frame_height))
|
||
|
||
if not out.isOpened():
|
||
print(f"无法创建视频文件: {output_path}")
|
||
return f"失败: {output_path}"
|
||
|
||
# 打开原视频（每个线程独立打开，避免线程安全问题）
|
||
cap = cv2.VideoCapture(input_video_path)
|
||
if not cap.isOpened():
|
||
out.release()
|
||
return f"失败: 无法打开视频 {input_video_path}"
|
||
|
||
# 生成视频帧
|
||
current_output_frame = 0
|
||
obj_frame_idx = 0
|
||
|
||
while current_output_frame < max_output_frame:
|
||
# 检查当前输出帧是否是5的倍数
|
||
if current_output_frame % frame_interval == 0 and obj_frame_idx < len(frame_list):
|
||
# 这是需要放置物体帧的位置
|
||
original_frame_id, xyxy = frame_list[obj_frame_idx]
|
||
|
||
# 设置原视频读取位置
|
||
cap.set(cv2.CAP_PROP_POS_FRAMES, original_frame_id)
|
||
ret, frame = cap.read()
|
||
|
||
if not ret:
|
||
# 读取失败，使用黑色帧
|
||
output_frame = np.zeros((frame_height, frame_width, 3), dtype=np.uint8)
|
||
else:
|
||
# 有数据，截取对应区域
|
||
x1, y1, x2, y2 = map(int, xyxy)
|
||
# 确保坐标在有效范围内
|
||
x1 = max(0, min(x1, frame_width))
|
||
y1 = max(0, min(y1, frame_height))
|
||
x2 = max(0, min(x2, frame_width))
|
||
y2 = max(0, min(y2, frame_height))
|
||
|
||
# 创建黑色背景
|
||
output_frame = np.zeros((frame_height, frame_width, 3), dtype=np.uint8)
|
||
# 将截取的区域放到输出帧中（保持原位置）
|
||
if x2 > x1 and y2 > y1:
|
||
cropped = frame[y1:y2, x1:x2]
|
||
output_frame[y1:y2, x1:x2] = cropped
|
||
|
||
# 移到下一个物体帧
|
||
obj_frame_idx += 1
|
||
else:
|
||
# 剩余帧留黑
|
||
output_frame = np.zeros((frame_height, frame_width, 3), dtype=np.uint8)
|
||
|
||
# 写入帧
|
||
out.write(output_frame)
|
||
current_output_frame += 1
|
||
|
||
# 释放资源
|
||
out.release()
|
||
cap.release()
|
||
print(f"已生成视频: {output_path}, 共 {current_output_frame} 帧")
|
||
return f"成功: {output_path}"
|
||
|
||
|
||
def frame_all_to_obj_vid(
|
||
json_data: dict,
|
||
input_video_path: str,
|
||
output_video_root: str,
|
||
) -> None:
|
||
"""
|
||
根据标注数据从原视频中截取物体，生成ai读取专用视频
|
||
|
||
参数:
|
||
json_data: 标注数据
|
||
input_video_path: 原视频路径
|
||
output_video_root: 输出视频根目录
|
||
"""
|
||
# 确保输出目录存在
|
||
os.makedirs(output_video_root, exist_ok=True)
|
||
|
||
# 1. 从 json_data 中提取数据，按 track_id 组织
|
||
track_dict: dict[int, list[tuple[int, list[int]]]] = {}
|
||
# 遍历每一帧
|
||
for frame_id_str, detections in json_data.items():
|
||
frame_id = int(frame_id_str)
|
||
for det in detections:
|
||
track_id = det.get("track_id", -1)
|
||
xyxy = det.get("xyxy", [0, 0, 0, 0])
|
||
if track_id not in track_dict:
|
||
track_dict[track_id] = []
|
||
track_dict[track_id].append((frame_id, xyxy))
|
||
|
||
# 2. 获取视频信息（只需要获取一次）
|
||
temp_cap = cv2.VideoCapture(input_video_path)
|
||
if not temp_cap.isOpened():
|
||
raise ValueError(f"无法打开视频: {input_video_path}")
|
||
|
||
frame_width = int(temp_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||
frame_height = int(temp_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||
temp_cap.release()
|
||
|
||
# 目标fps为25，每5帧取一帧（0, 5, 10...）
|
||
target_fps = 25
|
||
frame_interval = 5 # 每隔5帧取一帧
|
||
|
||
# 3. 使用多线程并行处理多个track_id
|
||
# 根据CPU核心数设置线程池大小
|
||
max_workers = min(os.cpu_count() or 4, len(track_dict))
|
||
print(f"使用 {max_workers} 个线程并行处理")
|
||
|
||
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||
# 提交所有任务
|
||
future_to_track = {
|
||
executor.submit(
|
||
process_track_id,
|
||
track_id,
|
||
frame_list,
|
||
input_video_path,
|
||
output_video_root,
|
||
frame_width,
|
||
frame_height,
|
||
target_fps,
|
||
frame_interval
|
||
):
|
||
track_id for track_id, frame_list in track_dict.items()
|
||
}
|
||
|
||
# 等待所有任务完成
|
||
for future in as_completed(future_to_track):
|
||
track_id = future_to_track[future]
|
||
try:
|
||
result = future.result()
|
||
print(f"Track ID {track_id}: {result}")
|
||
except Exception as e:
|
||
print(f"Track ID {track_id} 处理失败: {str(e)}")
|
||
|
||
# def frame_all_to_obj_vid(
|
||
# json_data: dict,
|
||
# input_video_path: str,
|
||
# output_video_root: str,
|
||
# ) -> None:
|
||
# """
|
||
# 根据标注数据从原视频中截取物体，生成ai读取专用视频
|
||
|
||
# 参数:
|
||
# json_data: 标注数据
|
||
# input_video_path: 原视频路径
|
||
# output_video_root: 输出视频根目录
|
||
# """
|
||
# # 确保输出目录存在
|
||
# os.makedirs(output_video_root, exist_ok=True)
|
||
|
||
# # 1. 从 json_data 中提取数据，按 track_id 组织
|
||
# track_dict: dict[int, list[tuple[int, list[int]]]] = {}
|
||
# # 遍历每一帧
|
||
# for frame_id_str, detections in json_data.items():
|
||
# frame_id = int(frame_id_str)
|
||
# for det in detections:
|
||
# track_id = det.get("track_id", -1)
|
||
# xyxy = det.get("xyxy", [0, 0, 0, 0])
|
||
# if track_id not in track_dict:
|
||
# track_dict[track_id] = []
|
||
# track_dict[track_id].append((frame_id, xyxy))
|
||
|
||
# # 2. 打开原视频
|
||
# cap = cv2.VideoCapture(input_video_path)
|
||
# if not cap.isOpened():
|
||
# raise ValueError(f"无法打开视频: {input_video_path}")
|
||
|
||
# # 获取原视频信息
|
||
# original_fps = cap.get(cv2.CAP_PROP_FPS)
|
||
# total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||
# frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||
# frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||
|
||
# # 目标fps为25，每5帧取一帧（0, 5, 10...）
|
||
# target_fps = 25
|
||
# frame_interval = 5 # 每隔5帧取一帧
|
||
|
||
# # 为每个 track_id 生成视频
|
||
# for track_id, frame_list in track_dict.items():
|
||
# # 计算需要生成的总帧数（确保覆盖所有物体帧且不少于两秒）
|
||
# min_frames_for_2s = 25 * 2 # 2秒 @ 25fps
|
||
# object_based_frames = len(frame_list) * frame_interval
|
||
# max_output_frame = max(object_based_frames, min_frames_for_2s)
|
||
|
||
# # 创建输出视频路径
|
||
# output_path = os.path.join(output_video_root, f"{track_id}.mp4")
|
||
|
||
# # 创建视频写入器
|
||
# fourcc = cv2.VideoWriter_fourcc(*'mp4v') # type: ignore
|
||
# out = cv2.VideoWriter(output_path, fourcc, target_fps, (frame_width, frame_height))
|
||
|
||
# if not out.isOpened():
|
||
# print(f"无法创建视频文件: {output_path}")
|
||
# continue
|
||
|
||
# # 生成视频帧
|
||
# current_output_frame = 0
|
||
# obj_frame_idx = 0
|
||
|
||
# while current_output_frame < max_output_frame:
|
||
# # 检查当前输出帧是否是5的倍数
|
||
# if current_output_frame % frame_interval == 0 and obj_frame_idx < len(frame_list):
|
||
# # 这是需要放置物体帧的位置
|
||
# original_frame_id, xyxy = frame_list[obj_frame_idx]
|
||
|
||
# # 设置原视频读取位置
|
||
# cap.set(cv2.CAP_PROP_POS_FRAMES, original_frame_id)
|
||
# ret, frame = cap.read()
|
||
|
||
# if not ret:
|
||
# # 读取失败，使用黑色帧
|
||
# output_frame = np.zeros((frame_height, frame_width, 3), dtype=np.uint8)
|
||
# else:
|
||
# # 有数据，截取对应区域
|
||
# x1, y1, x2, y2 = map(int, xyxy)
|
||
# # 确保坐标在有效范围内
|
||
# x1 = max(0, min(x1, frame_width))
|
||
# y1 = max(0, min(y1, frame_height))
|
||
# x2 = max(0, min(x2, frame_width))
|
||
# y2 = max(0, min(y2, frame_height))
|
||
|
||
# # 创建黑色背景
|
||
# output_frame = np.zeros((frame_height, frame_width, 3), dtype=np.uint8)
|
||
# # 将截取的区域放到输出帧中（保持原位置）
|
||
# if x2 > x1 and y2 > y1:
|
||
# cropped = frame[y1:y2, x1:x2]
|
||
# output_frame[y1:y2, x1:x2] = cropped
|
||
|
||
# # 移到下一个物体帧
|
||
# obj_frame_idx += 1
|
||
# else:
|
||
# # 剩余帧留黑
|
||
# output_frame = np.zeros((frame_height, frame_width, 3), dtype=np.uint8)
|
||
|
||
# # 写入帧
|
||
# out.write(output_frame)
|
||
# current_output_frame += 1
|
||
|
||
# # 释放视频写入器
|
||
# out.release()
|
||
# print(f"已生成视频: {output_path}, 共 {current_output_frame} 帧")
|
||
|
||
# # 释放原视频
|
||
# cap.release()
|
||
|
||
|
||
def create_mian_vid_for_ai(
|
||
input_video_path: str,
|
||
output_folder: str
|
||
) -> str:
|
||
"""
|
||
将原始视频的第0,1,2...帧映射到新视频的0,5,10...帧，其他帧留黑
|
||
|
||
参数:
|
||
input_video_path: 原始视频路径
|
||
output_folder: 输出文件夹路径
|
||
返回:
|
||
str: 输出视频路径
|
||
"""
|
||
# 确保输出目录存在
|
||
os.makedirs(output_folder, exist_ok=True)
|
||
|
||
# 构建输出视频路径
|
||
output_video_path = os.path.join(output_folder, "mian_vid_ai.mp4")
|
||
|
||
# 打开原视频
|
||
cap = cv2.VideoCapture(input_video_path)
|
||
if not cap.isOpened():
|
||
raise ValueError(f"无法打开视频: {input_video_path}")
|
||
|
||
# 获取原视频信息
|
||
original_fps = cap.get(cv2.CAP_PROP_FPS)
|
||
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||
|
||
# 目标fps为25，每5帧取一帧（0, 5, 10...）
|
||
target_fps = 25
|
||
frame_interval = 5
|
||
|
||
# 计算输出视频的总帧数
|
||
# 确保覆盖所有原始帧且不少于两秒
|
||
min_frames_for_2s = 25 * 2 # 2秒 @ 25fps
|
||
max_output_frame = max(total_frames * frame_interval, min_frames_for_2s)
|
||
|
||
# 创建视频写入器
|
||
fourcc = cv2.VideoWriter_fourcc(*'mp4v') # type: ignore
|
||
out = cv2.VideoWriter(output_video_path, fourcc, target_fps, (frame_width, frame_height))
|
||
|
||
if not out.isOpened():
|
||
raise ValueError(f"无法创建视频文件: {output_video_path}")
|
||
|
||
# 生成视频帧
|
||
current_output_frame = 0
|
||
original_frame_idx = 0
|
||
|
||
while current_output_frame < max_output_frame:
|
||
# 检查当前输出帧是否是5的倍数
|
||
if current_output_frame % frame_interval == 0 and original_frame_idx < total_frames:
|
||
# 这是需要放置原始帧的位置
|
||
# 设置原视频读取位置
|
||
cap.set(cv2.CAP_PROP_POS_FRAMES, original_frame_idx)
|
||
ret, frame = cap.read()
|
||
|
||
if not ret:
|
||
# 读取失败，使用黑色帧
|
||
output_frame = np.zeros((frame_height, frame_width, 3), dtype=np.uint8)
|
||
else:
|
||
# 有数据，使用原始帧
|
||
output_frame = frame
|
||
|
||
# 移到下一个原始帧
|
||
original_frame_idx += 1
|
||
else:
|
||
# 剩余帧留黑
|
||
output_frame = np.zeros((frame_height, frame_width, 3), dtype=np.uint8)
|
||
|
||
# 写入帧
|
||
out.write(output_frame)
|
||
current_output_frame += 1
|
||
|
||
# 释放资源
|
||
cap.release()
|
||
out.release()
|
||
print(f"已生成视频: {output_video_path}, 共 {current_output_frame} 帧")
|
||
|
||
return output_video_path |