87 lines
3.6 KiB
Python
87 lines
3.6 KiB
Python
import torch
|
|
|
|
class VideoFrameReorderAndSkip:
|
|
"""
|
|
一个用于处理视频图像序列的节点。
|
|
功能1：将图像序列对半切开，互换前后两半的顺序。
|
|
功能2：根据输入的数值进行抽帧（隔N取1）。
|
|
"""
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
# 输入端1：接收图像序列 (BATCH IMAGE)
|
|
"image_sequence": ("IMAGE",),
|
|
|
|
# 输入端2：抽帧数值填写 (INT)
|
|
# 含义：填写 N，表示每隔 N 帧保留一帧。
|
|
# 填写0 = 保留所有帧 (步长为1)
|
|
# 填写1 = 隔1抽1 (步长为2，保留第1、3、5...)
|
|
"skip_interval": ("INT", {
|
|
"default": 0, # 默认不抽帧
|
|
"min": 0, # 最小值不能小于0
|
|
"max": 120, # 设置一个合理的上限
|
|
"step": 1,
|
|
"display": "number"
|
|
}),
|
|
},
|
|
}
|
|
|
|
# 输出类型：处理后的图像序列
|
|
RETURN_TYPES = ("IMAGE",)
|
|
# 输出节点名称：在界面上显示的名称
|
|
RETURN_NAMES = ("processed_images",)
|
|
|
|
FUNCTION = "process_frames"
|
|
|
|
# 分类目录
|
|
CATEGORY = "My Custom Nodes/Video Tools"
|
|
|
|
def process_frames(self, image_sequence, skip_interval):
|
|
# --- 理解输入数据 ---
|
|
# 在 ComfyUI 中，image_sequence 是一个 PyTorch Tensor。
|
|
# 它的形状通常是 [B, H, W, C] 即 [帧数, 高度, 宽度, 通道数]。
|
|
# 我们只需要操作第0维 (B维度)。
|
|
|
|
total_frames = image_sequence.shape[0]
|
|
|
|
# 安全检查：如果帧数少于2，没法进行对半交换操作
|
|
if total_frames < 2:
|
|
# 如果只有1帧，直接根据抽帧逻辑返回（基本就是原样返回）
|
|
step_size = skip_interval + 1
|
|
return (image_sequence[::step_size], )
|
|
|
|
# --- 步骤 1: 把首尾帧拼接到中间 (对半互换) ---
|
|
# 例如 16帧：[0-7] 是前半, [8-15] 是后半。目标是变成 [8-15] 接 [0-7]
|
|
|
|
# 1. 计算中间点位置 (整除)
|
|
mid_point = total_frames // 2
|
|
|
|
# 2. 利用切片分离前后半段
|
|
# [:mid_point] 表示从开头取到中间点之前
|
|
first_half = image_sequence[:mid_point]
|
|
# [mid_point:] 表示从中间点取到最后
|
|
second_half = image_sequence[mid_point:]
|
|
|
|
# 3. 重新拼接 (Concatenation)
|
|
# torch.cat 将两个张量沿着指定维度 (dim=0 即帧维度) 连接起来
|
|
# 注意顺序：second_half 在前，first_half 在后
|
|
reordered_images = torch.cat((second_half, first_half), dim=0)
|
|
|
|
# --- 步骤 2: 抽帧处理 ---
|
|
# 需求：填写1，就是每隔1帧保存下来 (即取第 0, 2, 4 帧...)，这意味着步长(step)是2。
|
|
# 需求：填写2，就是每隔2帧保存下来 (即取第 0, 3, 6 帧...)，这意味着步长(step)是3。
|
|
# 结论：实际切片的步长 = 输入的间隔数值 + 1
|
|
|
|
step_size = skip_interval + 1
|
|
|
|
# 使用 Python/PyTorch 的切片语法 [start:end:step]
|
|
# ::step_size 表示从头到尾，每隔 step_size 取一个
|
|
final_images = reordered_images[::step_size]
|
|
|
|
# --- 返回结果 ---
|
|
# ComfyUI 要求必须返回一个元组，即使只有一个输出
|
|
return (final_images,) |