guyouyue-ComfyUI_VideoAutoSplit/video_extract_key_frames/service.py

# service.py   主要是功能逻辑
import math
import os
import subprocess

import cv2
import numpy as np
import torch


# ==================== 图像相似度计算相关 ====================

def calculate_similarity(img1, img2, method="SSIM"):
    """
    计算两张图片的相似度

    参数:
        img1, img2: 待比较的两个图像数组
        method: 相似度计算方法，可选 "SSIM"、"MSE"、"pHash"

    返回:
        float: 相似度分数，范围在 [0, 1] 之间
    """
    if method == "SSIM":
        return ssim_similarity_torch(img1, img2)
    elif method == "MSE":
        return mse_similarity_torch(img1, img2)
    elif method == "pHash":
        return phash_similarity(img1.numpy(), img2.numpy())
    return 0.0


def ssim_similarity(img1, img2, C1=0.01 ** 2, C2=0.03 ** 2):
    """
    使用结构相似性指数(SSIM)衡量两幅图像之间的相似程度。

    参数:
        img1, img2: 待比较的两个图像数组，形状应一致。
        C1, C2: SSIM 公式中的稳定常量，防止除零错误。

    返回:
        float: SSIM 相似度得分，范围在 [0, 1] 之间，越大越相似。
    """
    # 如果是彩色图像，先转为灰度
    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 = ((img1 - mu1) ** 2).mean()
    sigma2 = ((img2 - mu2) ** 2).mean()
    # 计算协方差
    sigma12 = ((img1 - mu1) * (img2 - mu2)).mean()
    # 计算SSIM
    ssim = ((2 * mu1 * mu2 + C1) * (2 * sigma12 + C2)) / ((mu1 ** 2 + mu2 ** 2 + C1) * (sigma1 + sigma2 + C2))
    # 保证结果在0~1之间
    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表示差异最大。

    参数:
        img1, img2: 待比较的两张图像。

    返回:
        float: 归一化的相似度分数。
    """
    mse_val = mse(img1, img2)
    # 归一化，假设像素值范围为0~1
    # mse=0时相似度为1，mse=1时相似度为0
    similarity = 1.0 - min(1.0, mse_val)
    return similarity


def phash_similarity(img1, img2, hash_size=8):
    """
    计算两张图片的pHash相似度（归一化为0~1，1为完全相同）

    参数:
        img1, img2: 待比较的两张图像。
        hash_size: 哈希尺寸，默认为 8。

    返回:
        float: 相似度分数，在 [0, 1] 范围内。
    """
    hash1 = phash(img1, hash_size)
    hash2 = phash(img2, hash_size)
    # 汉明距离
    dist = np.count_nonzero(hash1 != hash2)
    similarity = 1 - dist / (hash_size * hash_size)
    return similarity


def mse(img1, img2):
    """
    计算两张图片的均方误差（MSE）

    参数:
        img1, img2: 输入的两张图片，shape相同

    返回:
        mse: 均方误差，值越小越相似
    """
    # 如果是Tensor，先转为numpy
    if hasattr(img1, 'cpu') and hasattr(img1, 'numpy'):
        img1 = img1.numpy()
        # img1 = img1.cpu().numpy()
    if hasattr(img2, 'cpu') and hasattr(img2, '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]
        img2 = 0.299 * img2[..., 0] + 0.587 * img2[..., 1] + 0.114 * img2[..., 2]
    # 保证类型为float32
    img1 = img1.astype(np.float32)
    img2 = img2.astype(np.float32)
    # 计算MSE
    mse_val = np.mean((img1 - img2) ** 2)
    return mse_val


def phash(img, hash_size=8):
    """
    计算单个图像的感知哈希（pHash）

    参数:
        img: 输入的图像，shape为(H, W, C)或(H, W)
        hash_size: 哈希大小，默认8

    返回:
        hash: 长度为hash_size*hash_size的01向量
    """
    # 如果是Tensor，先转为numpy
    if hasattr(img, 'cpu') and hasattr(img, '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]
    img = img.astype(np.float32)
    # 缩放到32x32
    img = cv2.resize(img, (32, 32), interpolation=cv2.INTER_AREA)
    # 进行DCT变换
    dct = cv2.dct(img)
    # 取左上角hash_size*hash_size的DCT系数
    dct_low_freq = dct[:hash_size, :hash_size]
    # 计算中值
    med = np.median(dct_low_freq)
    # 生成hash
    hash = (dct_low_freq > med).astype(np.uint8)
    return hash.flatten()


# ==================== 关键帧提取相关 ====================

def compute_similarity_scores(images, similarity_method):
    """
    计算相邻帧之间的相似度分数

    参数:
        images: numpy数组格式的图像序列
        similarity_method: 相似度计算方法

    返回:
        list: 相似度分数列表
    """
    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):
    """
    基于统计方法计算相似度阈值

    参数:
        similarity_scores: 相似度分数列表
        threshold: 用户设定的阈值调节参数

    返回:
        float: 计算得到的相似度阈值
    """
    scores_array = np.array(similarity_scores)
    score_mean = scores_array.mean()
    score_std = scores_array.std()
    similarity_limit = score_mean - score_std * 1.5 - threshold

    # 旧版本算法
    # ssim_max = max(similarity_scores)
    # ssim_mean = sum(similarity_scores) / len(similarity_scores)
    # similarity_limit = ssim_max - (ssim_max - ssim_mean) * 2 - threshold

    return similarity_limit


def detect_keyframes_by_threshold(similarity_scores, similarity_limit):
    """
    根据阈值检测关键帧

    参数:
        similarity_scores: 相似度分数列表
        similarity_limit: 相似度阈值

    返回:
        list: 关键帧索引列表（总是包含第0帧）
    """
    keyframes = [0]  # 第0帧总是关键帧
    for i, score in enumerate(similarity_scores):
        if score < similarity_limit:
            keyframes.append(i + 1)
    return keyframes


def merge_close_keyframes(keyframes, total_frames, merge_inter_frames):
    """
    合并距离过近的关键帧

    参数:
        keyframes: 原始关键帧索引列表
        total_frames: 总帧数
        merge_inter_frames: 最小间隔帧数

    返回:
        list: 合并后的关键帧索引列表
    """
    if not keyframes:
        return []

    merged = [keyframes[0]]
    for kf in keyframes[1:]:
        if kf - merged[-1] > merge_inter_frames:
            merged.append(kf)

    # 优化尾部处理
    if merged and (total_frames - merged[-1]) < merge_inter_frames:
        # 尝试向前合并而不是循环删除
        if len(merged) > 1:
            merged.pop()

    return merged


def split_large_intervals(keyframes, total_frames, max_frames):
    """
    拆分超大间隔，确保单个片段不超过最大帧数

    参数:
        keyframes: 关键帧索引列表
        total_frames: 总帧数
        max_frames: 单个片段最大帧数

    返回:
        list: 拆分后的关键帧索引列表
    """
    result = []

    for i in range(len(keyframes)):
        start = keyframes[i]
        end = keyframes[i + 1] if i + 1 < len(keyframes) else total_frames

        result.append(start)

        interval_size = end - start
        if interval_size > max_frames:
            # 需要拆分
            num_splits = math.ceil(interval_size / max_frames)
            frames_per_split = interval_size // num_splits

            # 添加中间分割点
            for j in range(1, num_splits):
                split_point = start + j * frames_per_split
                if split_point < end:
                    result.append(split_point)

    # 移除重复和排序
    result = sorted(set(result))

    # 移除最后一个如果等于或超过total_frames（避免越界）
    if result and result[-1] >= total_frames:
        result = result[:-1]

    return result


def extract_keyframes(image, threshold, merge_inter_frames, max_frames, similarity_method):
    """
    从图像序列中提取关键帧的核心逻辑

    参数:
        image: 图像序列张量，维度为 [B, H, W, C]
        threshold: 控制关键帧检测灵敏度的浮点值
        merge_inter_frames: 最少合并间隔帧数
        max_frames: 单个批次最大帧数限制
        similarity_method: 选择使用的图像相似度算法

    返回:
        list: 关键帧索引列表
    """
    # 0. 基础检查
    if image.shape[0] < 1:
        raise ValueError("输入图像序列为空")

    total_frames = image.shape[0]

    # 边界情况：只有一帧
    if total_frames == 1:
        return [0]

    # 1. 计算相邻帧相似度
    images_np = image
    # images_np = image.cpu().numpy()
    similarity_scores = compute_similarity_scores(images_np, similarity_method)

    if not similarity_scores:
        return [0]

    # 2. 基于统计阈值检测关键帧
    similarity_limit = calculate_similarity_threshold(similarity_scores, threshold)
    keyframes = detect_keyframes_by_threshold(similarity_scores, similarity_limit)

    # 3. 合并过近的关键帧
    keyframes = merge_close_keyframes(keyframes, total_frames, merge_inter_frames)

    # 4. 拆分超大间隔
    keyframes = split_large_intervals(keyframes, total_frames, max_frames)

    return keyframes


# ==================== 视频处理相关 ====================

def get_ffmpeg_path():
    """
    获取ffmpeg可执行文件路径

    返回:
        str: ffmpeg路径

    异常:
        ProcessLookupError: 找不到ffmpeg时抛出
    """
    try:
        ffmpeg_path = subprocess.check_output(['which', 'ffmpeg'], text=True).strip()
        return ffmpeg_path
    except:
        # 尝试直接使用ffmpeg（假设在系统路径中）
        try:
            subprocess.run(['ffmpeg', '-version'], capture_output=True, check=True)
            return 'ffmpeg'
        except:
            raise ProcessLookupError("找不到 ffmpeg，请先安装。可以通过 apt install ffmpeg 或 brew install ffmpeg 安装")


def get_video_format_config(format):
    """
    获取视频格式配置

    参数:
        format: 视频格式名称

    返回:
        dict: 格式配置字典
    """
    format_configs = {
        "h264-mp4": {
            "ext": "mp4",
            "vcodec": "h264_nvenc" if torch.cuda.is_available() else "libx264",
            "pix_fmt": "yuv420p",
            "preset": "medium"
        },
        "h265-mp4": {
            "ext": "mp4",
            "vcodec": "hevc_nvenc" if torch.cuda.is_available() else "libx265",
            "pix_fmt": "yuv420p",
            "preset": "medium"
        },
        "vp9-webm": {
            "ext": "webm",
            "vcodec": "libvpx-vp9",
            "pix_fmt": "yuva420p",
            "preset": "good"
        }
    }
    return format_configs.get(format, format_configs["h264-mp4"])


def adjust_dimensions_for_codec(width, height):
    """
    调整尺寸以符合编码器要求（必须是偶数）

    参数:
        width: 原始宽度
        height: 原始高度

    返回:
        tuple: (调整后的宽度, 调整后的高度)
    """
    if width % 2 != 0:
        width -= 1
    if height % 2 != 0:
        height -= 1
    return width, height


def encode_video_segment(segment_images, output_path, frame_rate, config, video_quality, ffmpeg_path):
    """
    将图像序列编码为视频文件

    参数:
        segment_images: 图像序列张量
        output_path: 输出文件路径
        frame_rate: 帧率
        config: 视频格式配置
        video_quality: 视频质量（CRF值）
        ffmpeg_path: ffmpeg可执行文件路径

    返回:
        bool: 成功返回True，失败返回False
    """
    height, width = segment_images.shape[1], segment_images.shape[2]
    width, height = adjust_dimensions_for_codec(width, height)

    ffmpeg_args = [
        ffmpeg_path,
        '-y',  # 覆盖已存在的文件
        '-f', 'rawvideo',
        '-vcodec', 'rawvideo',
        '-s', f'{width}x{height}',
        '-pix_fmt', 'rgb24',
        '-r', str(frame_rate),
        '-i', '-',  # 从stdin读取
        '-c:v', config['vcodec'],
        '-crf', str(video_quality),
        '-preset', config['preset'],
        '-pix_fmt', config['pix_fmt'],
        '-color_range', 'pc',
        '-colorspace', 'bt709',
        '-color_primaries', 'bt709',
        '-color_trc', 'iec61966-2-1',
        output_path
    ]

    # 准备图像数据
    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, :]

    try:
        process = subprocess.Popen(
            ffmpeg_args,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE
        )

        for frame in images_np:
            process.stdin.write(frame.tobytes())

        process.stdin.close()
        process.wait()

        if process.returncode != 0:
            error_msg = process.stderr.read().decode('utf-8')
            print(f"FFmpeg错误: {error_msg}")
            return False

        return True

    except Exception as e:
        print(f"编码视频时出错: {str(e)}")
        return False


def mux_audio_video(video_path, audio_segment, audio_sample_rate, output_path, ffmpeg_path):
    """
    将音频和视频合并

    参数:
        video_path: 视频文件路径
        audio_segment: 音频数据张量
        audio_sample_rate: 音频采样率
        output_path: 输出文件路径
        ffmpeg_path: ffmpeg可执行文件路径

    返回:
        bool: 成功返回True，失败返回False
    """
    channels = audio_segment.size(1)
    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,
        '-y',
        '-i', video_path,
        '-ar', str(audio_sample_rate),
        '-ac', str(channels),
        '-f', 'f32le',
        '-i', '-',
        '-c:v', 'copy',
        '-c:a', 'aac',
        '-b:a', '192k',
        '-shortest',
        output_path
    ]

    try:
        process = subprocess.Popen(
            mux_args,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE
        )

        process.communicate(input=audio_data)

        if process.returncode != 0:
            return False

        return True

    except Exception as e:
        print(f"音频合并失败: {str(e)}")
        return False


def split_and_save_video_segments(image, keyframe_indices, frame_rate, format,
                                  video_quality, output_folder, filename, audio=None):
    """
    根据关键帧索引切割并保存视频片段

    参数:
        image: 输入图像序列张量
        keyframe_indices: 关键帧索引列表
        frame_rate: 输出视频帧率
        filename_prefix: 输出文件名前缀
        format: 视频编码格式
        video_quality: 视频质量（CRF值）
        output_folder: 输出文件夹路径
        filename: 基础文件名
        audio: 可选音频数据
        prompt: 元数据提示信息
        extra_pnginfo: 额外PNG信息

    返回:
        list: 保存的视频文件路径列表
    """
    # 获取ffmpeg路径和格式配置
    ffmpeg_path = get_ffmpeg_path()
    config = get_video_format_config(format)

    segment_files = []
    total_frames = image.shape[0]

    # 添加结束帧索引
    segment_indices = keyframe_indices + [total_frames]

    # 获取音频信息
    audio_waveform = None
    audio_sample_rate = None
    if audio is not None:
        try:
            audio_waveform = audio['waveform']
            audio_sample_rate = audio['sample_rate']
        except:
            pass

    print(f"\n开始切割视频，共 {len(segment_indices) - 1} 个片段")

    # 逐个片段处理
    for seg_idx in range(len(segment_indices) - 1):
        start_frame = segment_indices[seg_idx]
        end_frame = segment_indices[seg_idx + 1]

        segment_images = image[start_frame:end_frame]
        num_frames = segment_images.shape[0]

        if num_frames == 0:
            continue

        print(f"处理片段 {seg_idx + 1}: 帧 {start_frame}-{end_frame} (共{num_frames}帧)")

        # 生成文件名
        temp_video_file = f"{filename}_{seg_idx + 1:03d}_temp.{config['ext']}"
        temp_video_path = os.path.join(output_folder, temp_video_file)
        final_file = f"{filename}_{seg_idx + 1:03d}.{config['ext']}"
        final_path = os.path.join(output_folder, final_file)

        # 编码视频
        video_path = temp_video_path if audio_waveform is not None else final_path
        success = encode_video_segment(
            segment_images, video_path, frame_rate, config, video_quality, ffmpeg_path
        )

        if not success:
            print(f"✗ 片段 {seg_idx + 1} 编码失败")
            continue

        # 如果有音频，提取对应时间段的音频并合并
        if audio_waveform is not None:
            # 计算音频时间范围
            start_time = start_frame / frame_rate
            duration = num_frames / frame_rate

            # 提取音频片段
            start_sample = int(start_time * audio_sample_rate)
            end_sample = int((start_time + duration) * audio_sample_rate)

            # 确保索引不越界
            max_samples = audio_waveform.shape[2]
            start_sample = min(start_sample, max_samples)
            end_sample = min(end_sample, max_samples)

            if end_sample > start_sample:
                segment_audio = audio_waveform[:, :, start_sample:end_sample]

                # 尝试合并音频
                success = mux_audio_video(
                    temp_video_path, segment_audio, audio_sample_rate,
                    final_path, ffmpeg_path
                )

                if success:
                    # 删除临时文件
                    if os.path.exists(temp_video_path):
                        os.remove(temp_video_path)
                else:
                    # 使用无音频版本
                    print(f"音频合并失败，使用无音频版本")
                    if os.path.exists(temp_video_path):
                        os.rename(temp_video_path, final_path)
            else:
                # 音频片段无效，使用无音频版本
                print(f"音频片段无效，使用无音频版本")
                if os.path.exists(temp_video_path):
                    os.rename(temp_video_path, final_path)

        segment_files.append(final_path)
        print(f"✓ 片段 {seg_idx + 1} 已保存: {final_file}")

    return segment_files