diff --git a/essentials.py b/essentials.py index 8ab4857..388b176 100644 --- a/essentials.py +++ b/essentials.py @@ -4,11 +4,10 @@ import ast import math import random import operator as op -import numpy as np +#import numpy as np import torch import torch.nn.functional as F - import torchvision.transforms.v2 as T from nodes import MAX_RESOLUTION, SaveImage @@ -67,8 +66,19 @@ class ImageResize: def execute(self, image, width, height, keep_proportion, interpolation="nearest", condition="always"): if keep_proportion is True: _, oh, ow, _ = image.shape - width = ow if width == 0 else width - height = oh if height == 0 else height + + if width == 0 and oh < height: + width = MAX_RESOLUTION + elif width == 0 and oh >= height: + width = ow + + if height == 0 and ow < width: + height = MAX_RESOLUTION + elif height == 0 and ow >= width: + height = ow + + #width = ow if width == 0 else width + #height = oh if height == 0 else height ratio = min(width / ow, height / oh) width = round(ow*ratio) height = round(oh*ratio) @@ -913,6 +923,64 @@ class BatchCount: count = len(batch) return (count, ) + +class ImageSeamCarving: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "width": ("INT", { "default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1, }), + "height": ("INT", { "default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1, }), + "energy": (["backward", "forward"],), + "order": (["width-first", "height-first"],), + }, + "optional": { + "keep_mask": ("MASK",), + "drop_mask": ("MASK",), + } + } + + RETURN_TYPES = ("IMAGE",) + CATEGORY = "essentials" + FUNCTION = "execute" + + def execute(self, image, width, height, energy, order, keep_mask=None, drop_mask=None): + try: + from .carve import seam_carving + except ImportError as e: + raise Exception(e) + + img = p(image) + + if keep_mask is not None: + #keep_mask = keep_mask.reshape((-1, 1, keep_mask.shape[-2], keep_mask.shape[-1])).movedim(1, -1) + keep_mask = p(keep_mask.unsqueeze(-1)) + + if keep_mask.shape[2] != img.shape[2] or keep_mask.shape[3] != img.shape[3]: + keep_mask = F.interpolate(keep_mask, size=(img.shape[2], img.shape[3]), mode="bilinear") + if drop_mask is not None: + drop_mask = p(drop_mask.unsqueeze(-1)) + + if drop_mask.shape[2] != img.shape[2] or drop_mask.shape[3] != img.shape[3]: + drop_mask = F.interpolate(drop_mask, size=(img.shape[2], img.shape[3]), mode="bilinear") + + out = [] + for i in range(img.shape[0]): + resized = seam_carving( + T.ToPILImage()(img[i]), + size=(width, height), + energy_mode=energy, + order=order, + keep_mask=T.ToPILImage()(keep_mask[i]) if keep_mask is not None else None, + drop_mask=T.ToPILImage()(drop_mask[i]) if drop_mask is not None else None, + ) + out.append(T.ToTensor()(resized)) + + out = torch.stack(out) + out = pb(out) + + return(out, ) class CLIPTextEncodeSDXLSimplified: @classmethod @@ -974,6 +1042,7 @@ NODE_CLASS_MAPPINGS = { "ImageDesaturate+": ImageDesaturate, "ImagePosterize+": ImagePosterize, "ImageCASharpening+": ImageCAS, + "ImageSeamCarving+": ImageSeamCarving, "ImageEnhanceDifference+": ImageEnhanceDifference, "ImageExpandBatch+": ImageExpandBatch, "ImageFromBatch+": ImageFromBatch, @@ -1009,6 +1078,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "ImageDesaturate+": "🔧 Image Desaturate", "ImagePosterize+": "🔧 Image Posterize", "ImageCASharpening+": "🔧 Image Contrast Adaptive Sharpening", + "ImageSeamCarving+": "🔧 Image Seam Carving", "ImageEnhanceDifference+": "🔧 Image Enhance Difference", "ImageExpandBatch+": "🔧 Image Expand Batch", "ImageFromBatch+": "🔧 Image From Batch",