glh
f73c886569
feat(video):优化关键帧提取性能并支持GPU加速 - 引入torch库以支持GPU加速计算 - 实现SSIM和MSE的GPU版本相似度计算函数 - 修改图像相似度计算逻辑，优先使用GPU加速方法 - 调整关键帧合并策略，优化尾部关键帧处理逻辑 - 启用CUDA编码器（h264_nvenc/hevc_nvenc）以提升视频处理性能 - 移除不必要的隐藏参数以简化节点接口 - 添加处理时间打印功能用于性能监控- 优化Tensor到numpy转换方式，减少冗余操作
172 lines
6.9 KiB
Python
172 lines
6.9 KiB
Python
# nodes.py
|
|
# 主要是各个节点的定义和业务逻辑
|
|
from datetime import datetime
|
|
|
|
import folder_paths
|
|
|
|
from .service import (
|
|
extract_keyframes,
|
|
split_and_save_video_segments
|
|
)
|
|
|
|
|
|
class ExtractKeyFrames:
|
|
"""
|
|
用于从图像序列中提取关键帧的节点类。支持多种相似性计算方法（如 SSIM、MSE 和 pHash），
|
|
并根据设定阈值判断是否为显著变化的关键帧，并进行合并与拆分控制以适应后续处理需求。
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
"""
|
|
定义节点所需输入参数的信息。
|
|
|
|
返回:
|
|
dict: 包含 required 字段的字典，描述必填参数及属性：
|
|
- image: 输入的图像序列，格式为 IMAGE 类型。
|
|
- threshold: 浮点型，用于调节关键帧识别敏感度，默认 0.1。
|
|
- merge_inter_frames: 整数型，定义合并相邻关键帧的最小间隔，默认 10。
|
|
- max_frames: 整数型，单批次允许的最大帧数，默认 999999。
|
|
- similarity_method: 枚举型，选择使用的图像相似度算法，可选 ["SSIM", "MSE", "pHash"]。
|
|
"""
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"threshold": ("FLOAT", {"default": 0.1, "min": -1.0, "max": 1.00, "step": 0.01}),
|
|
"merge_inter_frames": ("INT", {"default": 10, "min": 0, "max": 999, "step": 1}),
|
|
"max_frames": ("INT", {"default": 999999, "min": 1, "max": 999999999, "step": 1}),
|
|
"similarity_method": (["SSIM", "MSE", "pHash"], {"default": "SSIM"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING", "INT")
|
|
RETURN_NAMES = ("KEYFRAMES", "indexes_string", "indexes_int")
|
|
|
|
CATEGORY = "ExtractKeyFrames(关键帧处理)"
|
|
FUNCTION = "execute"
|
|
|
|
def execute(self, image, threshold, merge_inter_frames, max_frames, similarity_method):
|
|
"""
|
|
执行关键帧提取流程。
|
|
|
|
参数:
|
|
image: 图像序列张量，维度为 [B, H, W, C]。
|
|
threshold: 控制关键帧检测灵敏度的浮点值。
|
|
merge_inter_frames: 最少合并间隔帧数。
|
|
max_frames: 单个批次最大帧数限制。
|
|
similarity_method: 选择使用的图像相似度算法。
|
|
|
|
返回:
|
|
tuple: 包含以下三个元素：
|
|
- 提取的关键帧图像张量。
|
|
- 关键帧索引组成的字符串。
|
|
- 关键帧索引组成的整数列表。
|
|
"""
|
|
# 调用service层的关键帧提取逻辑
|
|
start=datetime.now()
|
|
keyframe_indices = extract_keyframes(
|
|
image, threshold, merge_inter_frames, max_frames, similarity_method
|
|
)
|
|
|
|
# 生成返回结果
|
|
keyframe_images = image[keyframe_indices]
|
|
indexes_string = ','.join(map(str, keyframe_indices))
|
|
print(f"用时{datetime.now()-start}")
|
|
return (keyframe_images, indexes_string, keyframe_indices)
|
|
|
|
|
|
class ExtractKeyFramesAndSplit:
|
|
"""
|
|
根据关键帧提取结果，将原视频按关键帧位置切割成多个子视频片段。
|
|
支持音频同步切割和多种视频格式输出。
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"threshold": ("FLOAT", {"default": 0.1, "min": -1.0, "max": 1.00, "step": 0.01}),
|
|
"merge_inter_frames": ("INT", {"default": 10, "min": 0, "max": 999, "step": 1}),
|
|
"max_frames": ("INT", {"default": 999999, "min": 1, "max": 999999999, "step": 1}),
|
|
"similarity_method": (["SSIM", "MSE", "pHash"], {"default": "SSIM"}),
|
|
"frame_rate": ("INT", {"default": 8, "min": 1, "max": 120, "step": 1}),
|
|
"filename_prefix": ("STRING", {"default": "KeyFrameSegment"}),
|
|
"format": (["h264-mp4", "h265-mp4", "vp9-webm"], {"default": "h264-mp4"}),
|
|
"save_segments": ("BOOLEAN", {"default": True}),
|
|
"video_quality": ("INT", {"default": 23, "min": 0, "max": 51, "step": 1}),
|
|
},
|
|
"optional": {
|
|
"audio": ("AUDIO",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING", "INT", "STRING")
|
|
RETURN_NAMES = ("KEYFRAMES", "indexes_string", "indexes_int", "segment_files")
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "execute"
|
|
CATEGORY = "ExtractKeyFrames"
|
|
|
|
def execute(self, image, threshold, merge_inter_frames, max_frames,
|
|
similarity_method, frame_rate, filename_prefix, format,
|
|
save_segments=True, video_quality=23, audio=None,):
|
|
"""
|
|
执行关键帧提取和视频切割
|
|
|
|
参数:
|
|
image: 输入图像序列。
|
|
threshold: 相似度检测阈值。
|
|
merge_inter_frames: 合并相邻关键帧的最小间隔帧数。
|
|
max_frames: 单个片段最大帧数。
|
|
similarity_method: 相似度计算方法。
|
|
frame_rate: 输出视频帧率。
|
|
filename_prefix: 输出文件名前缀。
|
|
format: 视频编码格式。
|
|
save_segments: 是否保存分割后的视频片段。
|
|
video_quality: 视频质量（CRF值）。
|
|
audio: 可选音频数据。
|
|
prompt: 元数据提示信息。
|
|
extra_pnginfo: 额外PNG信息。
|
|
|
|
返回:
|
|
tuple: 包含关键帧图像、索引字符串、索引列表和视频片段路径列表。
|
|
"""
|
|
|
|
# 1. 提取关键帧索引（复用service层逻辑）
|
|
keyframe_indices = extract_keyframes(
|
|
image, threshold, merge_inter_frames, max_frames, similarity_method
|
|
)
|
|
|
|
# 2. 获取关键帧图像
|
|
keyframe_images = image[keyframe_indices]
|
|
indexes_string = ','.join(map(str, keyframe_indices))
|
|
|
|
# 3. 如果需要保存视频片段
|
|
segment_files = []
|
|
if save_segments:
|
|
# 获取输出路径
|
|
output_dir = folder_paths.get_output_directory()
|
|
full_output_folder, filename, _, subfolder, _ = folder_paths.get_save_image_path(
|
|
filename_prefix, output_dir
|
|
)
|
|
|
|
# 调用service层的视频切割逻辑（只传递必要参数）
|
|
segment_files = split_and_save_video_segments(
|
|
image=image,
|
|
keyframe_indices=keyframe_indices,
|
|
frame_rate=frame_rate,
|
|
format=format,
|
|
video_quality=video_quality,
|
|
output_folder=full_output_folder,
|
|
filename=filename,
|
|
audio=audio
|
|
)
|
|
|
|
segment_files_str = '\n'.join(segment_files) if segment_files else ""
|
|
|
|
print(f"\n✓ 完成！共提取 {len(keyframe_indices)} 个关键帧，生成 {len(segment_files)} 个视频片段")
|
|
|
|
return (keyframe_images, indexes_string, keyframe_indices, segment_files_str)
|
|
|
|
|