diff --git a/__init__.py b/__init__.py index bf4d75c..ecbef99 100644 --- a/__init__.py +++ b/__init__.py @@ -1,13 +1,26 @@ -from .nodes.MakeFrame import(BreakFrames, GetKeyFrames, MakeGrid) +import os +import sys +import subprocess + +here = os.path.dirname(__file__) +requirements_path = os.path.join(here, "requirements.txt") + +try: + from .nodes.MakeFrame import BreakFrames, GetKeyFrames, MakeGrid, BreakGrid +except: + subprocess.check_call([sys.executable, '-m', 'pip', 'install', '-r', requirements_path]) + from .nodes.MakeFrame import BreakFrames, GetKeyFrames, MakeGrid, BreakGrid NODE_CLASS_MAPPINGS = { "BreakFrames": BreakFrames, "GetKeyFrames": GetKeyFrames, "MakeGrid": MakeGrid, + "BreakGrid": BreakGrid, } NODE_DISPLAY_NAME_MAPPINGS = { "BreakFrames": "BreakFrames", "GetKeyFrames": "GetKeyFrames", "MakeGrid": "MakeGrid", + "BreakGrid": "BreakGrid", } \ No newline at end of file diff --git a/makeframeutils.py b/makeframeutils.py index bacb7b7..58e8db7 100644 --- a/makeframeutils.py +++ b/makeframeutils.py @@ -36,7 +36,6 @@ def normalize_size(images): if images[i].size != refimage.size: images[i] = images[i].resize(refimage.size, Image.Resampling.LANCZOS) return_images.append(images[i]) - np.lcm(6, 8) return return_images def constrain_image(image, max_width, max_height): @@ -187,8 +186,44 @@ def cat_to_pils(tensor): pils = [to_pil(tensor[i]) for i in range(tensor.shape[0])] return pils -def pil_to_cat(pimg): +def pil_to_tens(pimg): to_tensor = ToTensor() tensor = to_tensor(pimg).unsqueeze(0).permute(0, 2, 3, 1) - print(tensor.shape) - return tensor \ No newline at end of file + return tensor + +def get_grid_aspect(num_images: int, image_width: int, image_height: int) -> (int, int): + + if num_images == 0: + return 0, 0 + + min_diff = float('inf') + best_layout = (1, num_images) + + if image_width > image_height: + for cols in range(1, num_images + 1): + rows = -(-num_images // cols) + grid_width = cols * image_width + grid_height = rows * image_height + diff = abs(grid_width - grid_height) + + if diff < min_diff: + min_diff = diff + best_layout = (rows, cols) + + if cols > num_images / cols: + break + else: + for rows in range(1, num_images + 1): + cols = -(-num_images // rows) + grid_width = cols * image_width + grid_height = rows * image_height + diff = abs(grid_height - grid_width) + + if diff < min_diff: + min_diff = diff + best_layout = (rows, cols) + + if rows > num_images / rows: + break + + return best_layout \ No newline at end of file diff --git a/nodes/MakeFrame.py b/nodes/MakeFrame.py index 9d5a99a..8902122 100644 --- a/nodes/MakeFrame.py +++ b/nodes/MakeFrame.py @@ -1,6 +1,5 @@ import cv2 import torch -import torchvision.transforms as transforms import numpy as np import os from PIL import Image @@ -42,19 +41,16 @@ class BreakFrames: print(f"Error: Could not open video file {file_input}.") return None # Read each frame from anim - frame_count = 0 - transformer = transforms.ToTensor() tensors = [] while True: ret, frame = video_capture.read() if ret: frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) - tensors.append(transformer(frame).to(device).unsqueeze(0)) - frame_count += 1 + tensors.append(mfu.pil_to_tens(frame).to(device)) else: break video_capture.release() - cat_frame_tensors = torch.cat(tensors, dim = 0).permute(0, 2, 3, 1) + cat_frame_tensors = torch.cat(tensors, dim = 0).to(device) return (cat_frame_tensors,) @@ -75,8 +71,8 @@ class GetKeyFrames: }), }, } - RETURN_TYPES = ("IMAGE", "INT",) - RETURN_NAMES = ("Keyframes", "Keyframe indices",) + RETURN_TYPES = ("IMAGE", "IMAGE") + RETURN_NAMES = ("Keyframes", "Labeled Keyframes") FUNCTION = "getkeyframes" CATEGORY = "Frames" @@ -89,8 +85,14 @@ class GetKeyFrames: keyframe_indices = sorted([index.item() + 1 for index in top_indices]) keyframe_indices.insert(0, 0) cat_keyframe_tensors = [frames[i].to(device).unsqueeze(0) for i in keyframe_indices] - cat_keyframe_tensors = torch.cat(cat_keyframe_tensors, dim = 0) - return (cat_keyframe_tensors, keyframe_indices) + cat_keyframe_tensors = torch.cat(cat_keyframe_tensors, dim = 0).to(device) + + pils = mfu.cat_to_pils(cat_keyframe_tensors) + keyframes_labeled = [mfu.ImgLabeler(frame, str(idx), size=72, color="#ffffff") for frame, idx in zip(pils, keyframe_indices)] + keyframe_labeled_tensors = [mfu.pil_to_tens(keyframe) for keyframe in keyframes_labeled] + cat_keyframe_tensors_labeled = torch.cat(keyframe_labeled_tensors, dim = 0).to(device) + + return (cat_keyframe_tensors, cat_keyframe_tensors_labeled) class MakeGrid: def __init__(self): @@ -101,49 +103,67 @@ class MakeGrid: return { "required": { "frames": ("IMAGE", ), - "grid_rows": ("INT", { - "default": 4, - "min": 2, - "max": 24, - "step": 1 - }), - "grid_cols": ("INT", { - "default": 4, - "min": 2, - "max": 24, - "step": 1 - }), "max_width": ("INT", { - "default": 1024, + "default": 2048, "min": 64, - "max": 4096, + "max": 8000, "step": 8 }), "max_height": ("INT", { - "default": 1024, + "default": 2048, "min": 64, - "max": 4096, + "max": 8000, "step": 8 }), }, } - RETURN_TYPES = ("IMAGE", ) - RETURN_NAMES = ("Grid",) + RETURN_TYPES = ("IMAGE", "INT", "INT") + RETURN_NAMES = ("Grid", "Rows", "Columns") FUNCTION = "makegrid" CATEGORY = "Frames" - def makegrid(self, frames, grid_rows, grid_cols, max_width, max_height): - # Pad the list with extras; black space frustrates generation + def makegrid(self, frames, max_width, max_height): - # Size needs to be divisible by 8 AND the grid dimension - if not grid_rows == 8: max_height = mfu.closest_lcm(max_height, 8, grid_rows) - if not grid_cols == 8: max_width = mfu.closest_lcm(max_width, 8, grid_cols) - # Build base grid pils = mfu.cat_to_pils(frames) pils = mfu.normalize_size(pils) #normalize sizes to /8 - grid = mfu.constrain_image(mfu.MakeGrid(pils, grid_rows, grid_cols), max_width, max_height) - grid = mfu.pil_to_cat(grid) + rows, cols = mfu.get_grid_aspect(len(pils), pils[0].width, pils[0].height) + if len(pils) < rows*cols: + pils = mfu.padlist(pils, rows*cols) #pad list with repeats (black space bad) + if not rows == 8: max_height = mfu.closest_lcm(max_height, 8, rows) + if not cols == 8: max_width = mfu.closest_lcm(max_width, 8, cols) - return (grid,) \ No newline at end of file + grid = mfu.constrain_image(mfu.MakeGrid(pils, rows, cols), max_width, max_height) + grid = mfu.pil_to_tens(grid).to(device) + return (grid, rows, cols) + +class BreakGrid: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "grid": ("IMAGE",), + "rows": ("INT",{}), + "columns": ("INT",{}), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("Frames",) + + FUNCTION = "breakgrid" + CATEGORY = "Frames" + + def breakgrid(self, grid, rows, columns): + pilgrids = mfu.cat_to_pils(grid) + frames =[] + for pilgrid in pilgrids: + frames.extend(mfu.BreakGrid(pilgrid, rows, columns)) + 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 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..1db7aea --- /dev/null +++ b/requirements.txt @@ -0,0 +1 @@ +opencv-python \ No newline at end of file