diff --git a/__init__.py b/__init__.py index ecbef99..8e7f7c4 100644 --- a/__init__.py +++ b/__init__.py @@ -6,16 +6,17 @@ here = os.path.dirname(__file__) requirements_path = os.path.join(here, "requirements.txt") try: - from .nodes.MakeFrame import BreakFrames, GetKeyFrames, MakeGrid, BreakGrid + from .nodes.MakeFrame import BreakFrames, GetKeyFrames, MakeGrid, BreakGrid, RandomImageFromDir except: subprocess.check_call([sys.executable, '-m', 'pip', 'install', '-r', requirements_path]) - from .nodes.MakeFrame import BreakFrames, GetKeyFrames, MakeGrid, BreakGrid + from .nodes.MakeFrame import BreakFrames, GetKeyFrames, MakeGrid, BreakGrid, RandomImageFromDir NODE_CLASS_MAPPINGS = { "BreakFrames": BreakFrames, "GetKeyFrames": GetKeyFrames, "MakeGrid": MakeGrid, "BreakGrid": BreakGrid, + "RandomImageFromDir": RandomImageFromDir, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -23,4 +24,5 @@ NODE_DISPLAY_NAME_MAPPINGS = { "GetKeyFrames": "GetKeyFrames", "MakeGrid": "MakeGrid", "BreakGrid": "BreakGrid", + "RandomImageFromDir": "RandomImageFromDir", } \ No newline at end of file diff --git a/nodes/MakeFrame.py b/nodes/MakeFrame.py index 8951e85..e66aeef 100644 --- a/nodes/MakeFrame.py +++ b/nodes/MakeFrame.py @@ -2,9 +2,11 @@ import cv2 import torch import numpy as np import os +import random from PIL import Image from .. import makeframeutils as mfu + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") class BreakFrames: @@ -167,4 +169,39 @@ class BreakGrid: frame_tensors = [mfu.pil_to_tens(frame) for frame in frames] cat_frame_tensors = torch.cat(frame_tensors, dim = 0).unsqueeze(0) - return (cat_frame_tensors) \ No newline at end of file + return (cat_frame_tensors) + +class RandomImageFromDir: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "dir": ("STRING", { + "multiline": False, + "default": "C:/Poses" + }), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("Random Image",) + + FUNCTION = "getrandom" + CATEGORY = "Frames" + + def IS_CHANGED(cls, dir): + return random.random() + + def getrandom(self, dir): + files = os.listdir(dir) + image_files = [file for file in files if file.endswith(('.png', '.jpg', '.jpeg'))] + + random_image = random.choice(image_files) + + img = Image.open(os.path.join(dir, random_image)) + tensor = mfu.pil_to_tens(img).unsqueeze(0) + random_image = None + return tensor \ No newline at end of file