nodes/guyouyue-ComfyUI_VideoAutoSplit
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

v 0.0.3:

Browse Source 
feat(video):优化关键帧提取性能并支持GPU加速

- 引入torch库以支持GPU加速计算
- 实现SSIM和MSE的GPU版本相似度计算函数
- 修改图像相似度计算逻辑，优先使用GPU加速方法
- 调整关键帧合并策略，优化尾部关键帧处理逻辑
- 启用CUDA编码器（h264_nvenc/hevc_nvenc）以提升视频处理性能
- 移除不必要的隐藏参数以简化节点接口
- 添加处理时间打印功能用于性能监控- 优化Tensor到numpy转换方式，减少冗余操作
This commit is contained in:
glh
2025-10-21 00:08:14 +08:00
parent 8b4c2db25d
commit f73c886569
2 changed files with 88 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
video_extract_key_frames/nodes.py
+5 -8
View File
@@ -1,5 +1,7 @@
# nodes.py
# 主要是各个节点的定义和业务逻辑
from datetime import datetime
import folder_paths
from .service import (
@@ -61,6 +63,7 @@ class ExtractKeyFrames:
- 关键帧索引组成的整数列表。
"""
# 调用service层的关键帧提取逻辑
start=datetime.now()
keyframe_indices = extract_keyframes(
image, threshold, merge_inter_frames, max_frames, similarity_method
)
@@ -68,7 +71,7 @@ class ExtractKeyFrames:
# 生成返回结果
keyframe_images = image[keyframe_indices]
indexes_string = ','.join(map(str, keyframe_indices))
print(f"用时{datetime.now()-start}")
return (keyframe_images, indexes_string, keyframe_indices)
@@ -96,11 +99,6 @@ class ExtractKeyFramesAndSplit:
"optional": {
"audio": ("AUDIO",),
},
"hidden": {
"prompt": "PROMPT",
"extra_pnginfo": "EXTRA_PNGINFO",
"unique_id": "UNIQUE_ID"
},
}
RETURN_TYPES = ("IMAGE", "STRING", "INT", "STRING")
@@ -111,8 +109,7 @@ class ExtractKeyFramesAndSplit:
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,
prompt=None, extra_pnginfo=None):
save_segments=True, video_quality=23, audio=None,):
"""
执行关键帧提取和视频切割
video_extract_key_frames/service.py
+83 -35
View File
@@ -5,6 +5,7 @@ import subprocess
import cv2
import numpy as np
import torch
# ==================== 图像相似度计算相关 ====================
@@ -21,11 +22,11 @@ def calculate_similarity(img1, img2, method="SSIM"):
float: 相似度分数，范围在 [0, 1] 之间
"""
if method == "SSIM":
return ssim_similarity(img1, img2)
return ssim_similarity_torch(img1, img2)
elif method == "MSE":
return mse_similarity(img1, img2)
return mse_similarity_torch(img1, img2)
elif method == "pHash":
return phash_similarity(img1, img2)
return phash_similarity(img1.numpy(), img2.numpy())
return 0.0
@@ -58,6 +59,44 @@ def ssim_similarity(img1, img2, C1=0.01 ** 2, C2=0.03 ** 2):
return max(0.0, min(1.0, ssim.item() if hasattr(ssim, 'item') else ssim))
# 添加 GPU 加速的 SSIM 实现
def ssim_similarity_torch(img1, img2, C1=0.01 ** 2, C2=0.03 ** 2):
if img1.dim() == 4:
img1 = img1.squeeze(0)
if img2.dim() == 4:
img2 = img2.squeeze(0)
if img1.shape[-1] == 3:
img1 = 0.299 * img1[..., 0] + 0.587 * img1[..., 1] + 0.114 * img1[..., 2]
img2 = 0.299 * img2[..., 0] + 0.587 * img2[..., 1] + 0.114 * img2[..., 2]
mu1 = img1.mean()
mu2 = img2.mean()
sigma1 = torch.mean((img1 - mu1) ** 2)
sigma2 = torch.mean((img2 - mu2) ** 2)
sigma12 = torch.mean((img1 - mu1) * (img2 - mu2))
ssim = ((2 * mu1 * mu2 + C1) * (2 * sigma12 + C2)) / ((mu1 ** 2 + mu2 ** 2 + C1) * (sigma1 + sigma2 + C2))
return torch.clamp(ssim, 0.0, 1.0).item()
# 添加 GPU 加速的 MSE 实现
def mse_similarity_torch(img1, img2):
if img1.dim() == 4:
img1 = img1.squeeze(0)
if img2.dim() == 4:
img2 = img2.squeeze(0)
if img1.shape[-1] == 3:
img1 = 0.299 * img1[..., 0] + 0.587 * img1[..., 1] + 0.114 * img1[..., 2]
img2 = 0.299 * img2[..., 0] + 0.587 * img2[..., 1] + 0.114 * img2[..., 2]
mse_val = torch.mean((img1 - img2) ** 2)
similarity = 1.0 - torch.clamp(mse_val, 0.0, 1.0)
return similarity.item()
def mse_similarity(img1, img2):
"""
计算两张图片的MSE相似度，返回值归一化到0~1，1表示完全相同，0表示差异最大。
@@ -106,9 +145,11 @@ def mse(img1, img2):
"""
# 如果是Tensor，先转为numpy
if hasattr(img1, 'cpu') and hasattr(img1, 'numpy'):
img1 = img1.cpu().numpy()
img1 = img1.numpy()
# img1 = img1.cpu().numpy()
if hasattr(img2, 'cpu') and hasattr(img2, 'numpy'):
img2 = img2.cpu().numpy()
img2 = img2.numpy()
# img2 = img2.cpu().numpy()
# 如果是彩色图像，先转为灰度
if img1.ndim == 3 and img1.shape[-1] == 3:
img1 = 0.299 * img1[..., 0] + 0.587 * img1[..., 1] + 0.114 * img1[..., 2]
@@ -134,7 +175,8 @@ def phash(img, hash_size=8):
"""
# 如果是Tensor，先转为numpy
if hasattr(img, 'cpu') and hasattr(img, 'numpy'):
img = img.cpu().numpy()
img = img.numpy()
# img = img.cpu().numpy()
# 转为灰度
if img.ndim == 3 and img.shape[-1] == 3:
img = 0.299 * img[..., 0] + 0.587 * img[..., 1] + 0.114 * img[..., 2]
@@ -154,26 +196,35 @@ def phash(img, hash_size=8):
# ==================== 关键帧提取相关 ====================
def compute_similarity_scores(images_np, similarity_method):
def compute_similarity_scores(images, similarity_method):
"""
计算相邻帧之间的相似度分数
参数:
images_np: numpy数组格式的图像序列
images: numpy数组格式的图像序列
similarity_method: 相似度计算方法
返回:
list: 相似度分数列表
"""
similarity_scores = []
for i in range(len(images_np) - 1):
score = calculate_similarity(
images_np[i],
images_np[i + 1],
similarity_method
)
similarity_scores.append(score)
return similarity_scores
if similarity_method in ["SSIM", "MSE"]:
# 使用 GPU 加速的批量计算
return batch_compute_similarity(images, similarity_method)
else:
# pHash 仍然逐个计算
return [calculate_similarity(images[i], images[i + 1], similarity_method)
for i in range(len(images) - 1)]
def batch_compute_similarity(images, method):
# 准备成对的图像
img1_batch = images[:-1]
img2_batch = images[1:]
if method == "SSIM":
return [ssim_similarity_torch(img1, img2) for img1, img2 in zip(img1_batch, img2_batch)]
elif method == "MSE":
return [mse_similarity_torch(img1, img2) for img1, img2 in zip(img1_batch, img2_batch)]
def calculate_similarity_threshold(similarity_scores, threshold):
@@ -230,25 +281,19 @@ def merge_close_keyframes(keyframes, total_frames, merge_inter_frames):
返回:
list: 合并后的关键帧索引列表
"""
if len(keyframes) <= 1:
return keyframes
if not keyframes:
return []
# 从前往后合并，保留后一个
merged = [keyframes[0]]
for kf in keyframes[1:]:
if kf - merged[-1] > merge_inter_frames:
merged.append(kf)
else:
merged[-1] = kf # 替换为后一个
# 处理尾部：如果最后一个关键帧距离结尾太近，向前归并
max_iterations = len(merged) - 1
iteration = 0
while (len(merged) > 1 and
(total_frames - merged[-1]) < merge_inter_frames and
iteration < max_iterations):
merged.pop()
iteration += 1
# 优化尾部处理
if merged and (total_frames - merged[-1]) < merge_inter_frames:
# 尝试向前合并而不是循环删除
if len(merged) > 1:
merged.pop()
return merged
@@ -320,7 +365,8 @@ def extract_keyframes(image, threshold, merge_inter_frames, max_frames, similari
return [0]
# 1. 计算相邻帧相似度
images_np = image.cpu().numpy()
images_np = image
# images_np = image.cpu().numpy()
similarity_scores = compute_similarity_scores(images_np, similarity_method)
if not similarity_scores:
@@ -376,13 +422,13 @@ def get_video_format_config(format):
format_configs = {
"h264-mp4": {
"ext": "mp4",
"vcodec": "libx264",
"vcodec": "h264_nvenc" if torch.cuda.is_available() else "libx264",
"pix_fmt": "yuv420p",
"preset": "medium"
},
"h265-mp4": {
"ext": "mp4",
"vcodec": "libx265",
"vcodec": "hevc_nvenc" if torch.cuda.is_available() else "libx265",
"pix_fmt": "yuv420p",
"preset": "medium"
},
@@ -453,7 +499,8 @@ def encode_video_segment(segment_images, output_path, frame_rate, config, video_
]
# 准备图像数据
images_np = (segment_images.cpu().numpy() * 255).astype(np.uint8)
images_np = (segment_images.numpy() * 255).astype(np.uint8)
# images_np = (segment_images.cpu().numpy() * 255).astype(np.uint8)
if width != segment_images.shape[2] or height != segment_images.shape[1]:
# 裁剪到合适尺寸
images_np = images_np[:, :height, :width, :]
@@ -499,7 +546,8 @@ def mux_audio_video(video_path, audio_segment, audio_sample_rate, output_path, f
bool: 成功返回True，失败返回False
"""
channels = audio_segment.size(1)
audio_data = audio_segment.squeeze(0).transpose(0, 1).cpu().numpy().tobytes()
audio_data = audio_segment.squeeze(0).transpose(0, 1).numpy().tobytes()
# audio_data = audio_segment.squeeze(0).transpose(0, 1).cpu().numpy().tobytes()
mux_args = [
ffmpeg_path,