112 lines
3.6 KiB
Python
112 lines
3.6 KiB
Python
import cv2
|
|
import torch
|
|
import numpy as np
|
|
|
|
class SplitVideoByFrames:
|
|
"""
|
|
传入视频路径，根据帧率将视频拆分为多段，每段帧数上限可控，输出为张量序列列表
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"video_path": ("STRING", {"default": "your_video.mp4"}),
|
|
"max_frames_per_clip": ("INT", {"default": 30, "min": 1, "max": 1000}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LIST", "INT")
|
|
RETURN_NAMES = ("clips", "num_clips")
|
|
FUNCTION = "split_video"
|
|
CATEGORY = "云服务"
|
|
|
|
def split_video(self, video_path, max_frames_per_clip):
|
|
clips = []
|
|
cap = cv2.VideoCapture(video_path)
|
|
if not cap.isOpened():
|
|
raise FileNotFoundError(f"无法打开视频文件: {video_path}")
|
|
frames = []
|
|
while True:
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
break
|
|
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
|
np_image = frame.astype(np.float32) / 255.0
|
|
frames.append(torch.from_numpy(np_image))
|
|
if len(frames) == max_frames_per_clip:
|
|
clips.append(torch.stack(frames, dim=0))
|
|
frames = []
|
|
if frames:
|
|
clips.append(torch.stack(frames, dim=0))
|
|
cap.release()
|
|
return (clips, len(clips))
|
|
|
|
class GetVideoClipByIndex:
|
|
"""
|
|
传入 clips（张量序列列表）和索引，返回该片段合成的视频（mp4二进制）
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"clips": ("LIST", {}),
|
|
"index": ("INT", {"default": 0, "min": 0}),
|
|
"fps": ("INT", {"default": 24, "min": 1, "max": 120}),
|
|
"ext": (["mp4", "avi", "mov"], {"default": "mp4"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("BYTES",)
|
|
RETURN_NAMES = ("video_bytes",)
|
|
FUNCTION = "get_clip_video"
|
|
CATEGORY = "云服务"
|
|
|
|
def get_clip_video(self, clips, index, fps, ext):
|
|
import tempfile
|
|
import imageio_ffmpeg
|
|
import os
|
|
import io
|
|
if not clips or index < 0 or index >= len(clips):
|
|
raise IndexError("索引超出clips范围")
|
|
clip = clips[index]
|
|
arr = clip.cpu().numpy() if hasattr(clip, 'cpu') else clip
|
|
img_list = [np.clip(255. * img, 0, 255).astype(np.uint8) for img in arr]
|
|
height, width = img_list[0].shape[0], img_list[0].shape[1]
|
|
pix_fmt = 'rgb24'
|
|
with tempfile.NamedTemporaryFile(suffix=f'.{ext}', delete=False) as tmpfile:
|
|
tmp_path = tmpfile.name
|
|
cmd = [
|
|
imageio_ffmpeg.get_ffmpeg_exe(),
|
|
'-y',
|
|
'-f', 'rawvideo',
|
|
'-vcodec', 'rawvideo',
|
|
'-s', f'{width}x{height}',
|
|
'-pix_fmt', pix_fmt,
|
|
'-r', str(fps),
|
|
'-i', '-',
|
|
'-an',
|
|
'-vcodec', 'libx264',
|
|
'-pix_fmt', 'yuv420p',
|
|
tmp_path
|
|
]
|
|
import subprocess
|
|
proc = subprocess.Popen(cmd, stdin=subprocess.PIPE)
|
|
for im in img_list:
|
|
proc.stdin.write(im.tobytes())
|
|
proc.stdin.close()
|
|
proc.wait()
|
|
with open(tmp_path, "rb") as f:
|
|
video_bytes = f.read()
|
|
os.remove(tmp_path)
|
|
return (video_bytes,)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"SplitVideoByFrames": SplitVideoByFrames,
|
|
"GetVideoClipByIndex": GetVideoClipByIndex
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"SplitVideoByFrames": "🎬 视频分段拆帧",
|
|
"GetVideoClipByIndex": "🎞️ 获取片段视频"
|
|
}
|