diff --git a/__init__.py b/__init__.py index 57078ae..7fd117c 100644 --- a/__init__.py +++ b/__init__.py @@ -5,7 +5,7 @@ from PIL import Image import numpy as np import torch import comfy.utils -from .videoCut import getCutList, videoToPng, cutToDir +from .videoCut import getCutList, videoToPng, cutToDir, frames_to_video from .seg import get_masks from .thick_lines_from_canny import fill_white_segments, find_largest_white_component @@ -448,14 +448,14 @@ class videoCut: "video_path": ("STRING", {"default": "Temp"}), "save_name": ("STRING", {"default": ""}), "frame_rate": ("INT", { - "default": 1, + "default": 24, "min": 1, "max": 4096, "step": 1, "display": "number" }), "min_frame": ("INT", { - "default": 1, + "default": 16, "min": 1, "max": 4096, "step": 1, @@ -485,6 +485,42 @@ class videoCut: return (dirPathString,) +class frameToVideo: + + def __init__(self) -> None: + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "frame_dir": ("STRING", {"default": ""}), + "save_name": ("STRING", {"default": "result.mp4"}), + "frame_rate": ("INT", { + "default": 24, + "min": 1, + "max": 4096, + "step": 1, + "display": "number" + }), + } + } + + RETURN_TYPES = ("STRING",) + FUNCTION = "frame_to_video" + + CATEGORY = "badger" + + def frame_to_video(self, frame_dir, save_name, frame_rate): + output_abs = os.path.abspath(frame_dir) + output_dir = os.path.dirname(output_abs) + output_name = save_name + output_path = os.path.join(output_dir, output_name) + frames_to_video(frame_dir, frame_rate, output_path) + + return (output_path,) + + class getParentDir: def __init__(self) -> None: @@ -912,6 +948,7 @@ NODE_CLASS_MAPPINGS = { "FindThickLinesFromCanny-badger": FindThickLinesFromCanny, "TrimTransparentEdges-badger": TrimTransparentEdges, "ExpandImageWithColor-badger": ExpandImageWithColor, + "frameToVideo-badger": frameToVideo, } NODE_DISPLAY_NAME_MAPPINGS = { diff --git a/videoCut.py b/videoCut.py index 1342a47..42f49a2 100644 --- a/videoCut.py +++ b/videoCut.py @@ -104,6 +104,46 @@ def videoToPng(videopath, rate, save_name): return root_dir +def frames_to_video(input_folder, frame_rate, output_path): + # Get a sorted list of (non-hidden) image files + files = [f for f in sorted(os.listdir(input_folder)) if not f.startswith('.')] + sorted_files = sorted(files, key=lambda x: os.path.splitext(x)[0]) + + # Create a temporary file listing all images for ffmpeg + temp_list_path = 'ffmpeg_temp_list.txt' + with open(temp_list_path, 'w') as filelist: + for filename in sorted_files[:-1]: + file_path = os.path.join(input_folder, filename) + filelist.write(f"file '{file_path}'\n") + filelist.write(f"duration {1 / frame_rate}\n") + # Write the last file without a duration + filelist.write(f"file '{os.path.join(input_folder, sorted_files[-1])}'\n") + + # Construct the ffmpeg command to create the video + command = [ + 'ffmpeg', + '-f', 'concat', # Use the concat demuxer + '-safe', '0', # Allow unsafe file paths + '-i', temp_list_path, # Input file list + '-vsync', 'vfr', # Variable frame rate mode to match input sequence + '-pix_fmt', 'yuv420p', # Pixel format, widely compatible + '-c:v', 'libx264', # Codec to use for encoding + '-crf', '23', # Constant Rate Factor, controls quality (lower is better) + output_path + ] + + # Execute the command + try: + subprocess.run(command, check=True) + print(f'Video has been created at {output_path}') + except subprocess.CalledProcessError as e: + print(f'An error occurred: {e}') + finally: + # Clean up temporary file + if os.path.exists(temp_list_path): + os.remove(temp_list_path) + + def cutToDir(root_dir, cutList): dirIndex = 0 pngList = os.listdir(root_dir)