ETK Video based nodes!
This commit is contained in:
Tasha Upchurch
+2
-5
@@ -24,14 +24,9 @@ from custom_nodes.EternalKernelLiteGraphNodes.image import SelectFromRGBSimilari
|
||||
from custom_nodes.EternalKernelLiteGraphNodes.image import LoadImage
|
||||
from custom_nodes.EternalKernelLiteGraphNodes.image import PromptTemplate
|
||||
from custom_nodes.EternalKernelLiteGraphNodes.image import ExecWidget
|
||||
|
||||
from custom_nodes.EternalKernelLiteGraphNodes.image import ImageDistanceMask
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
if use_generic_nodes:
|
||||
NODE_CLASS_MAPPINGS_GENERIC = {
|
||||
"Interpolation": LatentInterpolation,
|
||||
@@ -76,3 +71,5 @@ else:
|
||||
NODE_CLASS_MAPPINGS.update(NODE_CLASS_MAPPINGS_COMFY)
|
||||
|
||||
|
||||
from custom_nodes.EternalKernelLiteGraphNodes.video import NODE_CLASS_MAPPINGS as NODE_CLASS_MAPPINGS_VIDEO
|
||||
NODE_CLASS_MAPPINGS.update(NODE_CLASS_MAPPINGS_VIDEO)
|
||||
@@ -0,0 +1,238 @@
|
||||
### info for code completion AI ###
|
||||
"""
|
||||
all of these classes are plugins for comfyui and follow the same pattern
|
||||
all of the images are torch tensors and it is unknown and unimportant if they are on the cpu or gpu
|
||||
all image inputs are (B,W,H,C)
|
||||
|
||||
avoid numpy and PIL as much as possible
|
||||
"""
|
||||
import os
|
||||
from abc import ABC, ABCMeta, abstractmethod
|
||||
from typing import List, Tuple, Union
|
||||
import cv2
|
||||
import torch
|
||||
from cv2 import VideoWriter, VideoWriter_fourcc, VideoCapture
|
||||
from torch import dtype
|
||||
|
||||
NODE_CLASS_MAPPINGS = {} # this is the dictionary that will be used to register the nodes
|
||||
|
||||
# setup directories
|
||||
root_p = os.path.dirname(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))
|
||||
input_dir = os.path.join(root_p, 'input', 'video')
|
||||
output_dir = os.path.join(root_p, 'output', 'video')
|
||||
|
||||
if not os.path.exists(input_dir):
|
||||
os.makedirs(input_dir)
|
||||
if not os.path.exists(output_dir):
|
||||
os.makedirs(output_dir)
|
||||
|
||||
out_video_folders = lambda s: sorted(
|
||||
[name for name in os.listdir(s.output_dir) if os.path.isdir(os.path.join(s.output_dir, name))])
|
||||
|
||||
in_video_folders = lambda s: sorted(
|
||||
[name for name in os.listdir(s.input_dir) if os.path.isdir(os.path.join(s.input_dir, name))])
|
||||
|
||||
# any files in in_video_folders ending in .mp4
|
||||
in_video_files = lambda s: sorted(
|
||||
[name for name in os.listdir(s.input_dir) if
|
||||
os.path.isfile(os.path.join(s.input_dir, name)) and name.endswith('.mp4')])
|
||||
|
||||
|
||||
# start with an ABC to define the common widget interface
|
||||
|
||||
class WidgetMetaclass(ABCMeta):
|
||||
"""A metaclass that automatically registers classes."""
|
||||
|
||||
def __init__(cls, name, bases, attrs):
|
||||
if (ABC in bases) or (VideoWidget in bases):
|
||||
pass
|
||||
else:
|
||||
NODE_CLASS_MAPPINGS[name] = cls
|
||||
|
||||
super().__init__(name, bases, attrs)
|
||||
|
||||
|
||||
class VideoWidget(ABC, metaclass=WidgetMetaclass):
|
||||
"""Abstract base class for simple construction"""
|
||||
|
||||
input_dir = input_dir
|
||||
output_dir = output_dir
|
||||
|
||||
CATEGORY = "video"
|
||||
RETURN_TYPES = ("VIDEO",)
|
||||
FUNCTION = "handler"
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def INPUT_TYPES(cls):
|
||||
"""All subclasses must provide a INPUT_TYPES classmethod"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def handler(self, one: object, two: object):
|
||||
"""All subclasses must provide a handler method"""
|
||||
pass
|
||||
|
||||
|
||||
class VideoFileToVideoFolder(VideoWidget, metaclass=WidgetMetaclass):
|
||||
"""Abstract base class for simple construction"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required":
|
||||
{
|
||||
"video_in": (in_video_files(cls),),
|
||||
"folder_out": (out_video_folders(cls),),
|
||||
},
|
||||
|
||||
"optional":
|
||||
{"text":
|
||||
("STRING", {"multiline": False}),
|
||||
}
|
||||
}
|
||||
|
||||
@abstractmethod
|
||||
def handler(self, one, two, text):
|
||||
return (one, two,)
|
||||
|
||||
class VideoFolderToImage(VideoWidget, metaclass=WidgetMetaclass):
|
||||
"""Abstract base class for simple construction"""
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required":
|
||||
{
|
||||
"folder_in": (in_video_folders(cls),),
|
||||
"idx_slice_start": ("INT", {"min": 0, "max": 10000, "step": 1, "default": 0}),
|
||||
"idx_slice_stop": ("INT", {"min": 0, "max": 10000, "step": 1, "default": 1}),
|
||||
"slice_idx_step": ("INT", {"min": 0, "max": 1000, "step": 1, "default": 1}),
|
||||
"idx": ("INT", {"min": 0, "max": 10000, "step": 1, "default": 0}),
|
||||
},
|
||||
|
||||
"optional":
|
||||
{"text":
|
||||
("STRING", {"multiline": False}),
|
||||
}
|
||||
}
|
||||
|
||||
@abstractmethod
|
||||
def handler(self, one, two, text):
|
||||
return (one,)
|
||||
|
||||
class VideoFileToImage(VideoWidget, metaclass=WidgetMetaclass):
|
||||
"""Abstract base class for simple construction"""
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required":
|
||||
{
|
||||
"video_in": (in_video_files(cls),),
|
||||
"idx_slice_start": ("INT", {"min": 0, "max": 10000, "step": 1, "default": 0}),
|
||||
"idx_slice_stop": ("INT", {"min": 0, "max": 10000, "step": 1, "default": 1}),
|
||||
"slice_idx_step": ("INT", {"min": 0, "max": 1000, "step": 1, "default": 1}),
|
||||
"idx": ("INT", {"min": 0, "max": 10000, "step": 1, "default": 0}),
|
||||
},
|
||||
|
||||
"optional":
|
||||
{"text":
|
||||
("STRING", {"multiline": False}),
|
||||
}
|
||||
}
|
||||
|
||||
@abstractmethod
|
||||
def handler(self, one, two, text):
|
||||
return (one,)
|
||||
|
||||
## end of abstract base classes ##
|
||||
|
||||
class VideoToFrames(VideoFileToVideoFolder, metaclass=WidgetMetaclass):
|
||||
"""use cv2 to open the video and use it as a source for image stacks"""
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.old = self.INPUT_TYPES
|
||||
# self.INPUT_TYPES = lambda s: self.old().update({})
|
||||
|
||||
def handler(self, video_in, folder_out, text):
|
||||
video_in_path = os.path.join(self.input_dir, video_in)
|
||||
video_out_path = os.path.join(self.output_dir, folder_out)
|
||||
|
||||
vidcap = cv2.VideoCapture(video_in_path)
|
||||
success, image = vidcap.read()
|
||||
count = 0
|
||||
while success:
|
||||
cv2.imwrite(os.path.join(video_out_path, "frame%d.jpg" % count), image)
|
||||
success, image = vidcap.read()
|
||||
count += 1
|
||||
|
||||
return (folder_out,)
|
||||
|
||||
|
||||
class VideoFramesToImageStack(VideoFolderToImage, metaclass=WidgetMetaclass):
|
||||
"""Use cv2 to open the video and save the videos individual frames"""
|
||||
|
||||
def handler(self, folder_in, text, idx_slice_start, idx_slice_stop, slice_idx_step, idx):
|
||||
"""
|
||||
Return slices of time from the video, sequences of images, defined by the start, stop and step
|
||||
- torch tensors (T,H,W,C)
|
||||
"""
|
||||
folder_in_path = os.path.join(self.input_dir, folder_in)
|
||||
image_list = []
|
||||
|
||||
for i in range(idx_slice_start+idx, idx_slice_stop+idx, slice_idx_step):
|
||||
file_fps = [os.path.join(folder_in_path, f"frame{i+ii}.jpg") for ii in range(slice_idx_step)]
|
||||
for fn in file_fps:
|
||||
image = cv2.imread(fn)
|
||||
if image is not None:
|
||||
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # converting to RGB format
|
||||
image_list.append(image)
|
||||
|
||||
# Convert list of images to tensor (T,H,W,C)
|
||||
tensor = torch.stack([torch.tensor(img,dtype=torch.float32) for img in image_list])
|
||||
tensor=tensor/255.0
|
||||
|
||||
|
||||
return (tensor,)
|
||||
|
||||
|
||||
class VideoFileToImageStack(VideoFileToImage, metaclass=WidgetMetaclass):
|
||||
"""Use cv2 to open a video and save the videos individual frames"""
|
||||
|
||||
def handler(self, video_in, text, idx_slice_start, idx_slice_stop, slice_idx_step, idx):
|
||||
"""
|
||||
Return slices of time from the video, sequences of images, defined by the start, stop and step
|
||||
- torch tensors (T,H,W,C)
|
||||
"""
|
||||
image_list = []
|
||||
video_in_fp = os.path.join(self.input_dir, video_in)
|
||||
vidcap = cv2.VideoCapture(video_in_fp)
|
||||
if not vidcap.isOpened():
|
||||
print("Error: Unable to open video file")
|
||||
return
|
||||
|
||||
# Seek to idx_slice_start+idx frame
|
||||
vidcap.set(cv2.CAP_PROP_POS_FRAMES, idx_slice_start + idx)
|
||||
|
||||
count = idx_slice_start + idx
|
||||
while count < idx_slice_stop + idx:
|
||||
success, image = vidcap.read()
|
||||
if not success:
|
||||
print("Error: Unable to read frame from video file")
|
||||
break
|
||||
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
||||
image_list.append(image)
|
||||
count += slice_idx_step
|
||||
# seek to next frame
|
||||
vidcap.set(cv2.CAP_PROP_POS_FRAMES, count)
|
||||
|
||||
# Convert list of images to tensor (T,H,W,C)
|
||||
tensor = torch.stack([torch.tensor(img, dtype=torch.float32) for img in image_list])
|
||||
tensor = tensor / 255.0
|
||||
|
||||
return (tensor,)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(NODE_CLASS_MAPPINGS)
|
||||
Reference in New Issue
Block a user