diff --git a/video_extract_key_frames/nodes.py b/video_extract_key_frames/nodes.py index 39fccac..08bb52c 100644 --- a/video_extract_key_frames/nodes.py +++ b/video_extract_key_frames/nodes.py @@ -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,): """ 执行关键帧提取和视频切割 diff --git a/video_extract_key_frames/service.py b/video_extract_key_frames/service.py index 516558f..e642827 100644 --- a/video_extract_key_frames/service.py +++ b/video_extract_key_frames/service.py @@ -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,