import torch import math import numpy as np from PIL import Image, ImageOps, ImageDraw from .util import tensor2pil, pil2tensor, empty_pil_tensor, crop_image import cv2 class OFFImageResizeFit: def __init__(self): pass @classmethod def INPUT_TYPES(s): return { "required":{ "image":("IMAGE",) }, "optional":{ "size":("INT",{ "default": 512, }) } } RETURN_TYPES = ("IMAGE",) FUNCTION = "doit" CATEGORY = "OFF" def doit(self, image, size=512): image = tensor2pil(image) new_width = size new_height = size ratio = float(image.width) / float(image.height) if image.width > image.height : new_width = int(size* ratio) else: new_height = int(size/ratio) image = image.resize((new_width, new_height)) return (pil2tensor(image),) class OFFCenterCrop: def __init__(self): pass @classmethod def INPUT_TYPES(s): """ Return a dictionary which contains config for all input fields. Some types (string): "MODEL", "VAE", "CLIP", "CONDITIONING", "LATENT", "IMAGE", "INT", "STRING", "FLOAT". Input types "INT", "STRING" or "FLOAT" are special values for fields on the node. The type can be a list for selection. Returns: `dict`: - Key input_fields_group (`string`): Can be either required, hidden or optional. A node class must have property `required` - Value input_fields (`dict`): Contains input fields config: * Key field_name (`string`): Name of a entry-point method's argument * Value field_config (`tuple`): + First value is a string indicate the type of field or a list for selection. + Secound value is a config for type "INT", "STRING" or "FLOAT". """ return { "required": { "image": ("IMAGE",) }, } RETURN_TYPES = ("IMAGE",) # RETURN_NAMES = ("image_output_name",) FUNCTION = "image_crop" # OUTPUT_NODE = False CATEGORY = "OFF" # Tensor to PIL def image_crop(self, image): image = tensor2pil(image) img_width, img_height = image.size crop_size = img_width if img_width > img_height: crop_size = img_height top = (img_height - crop_size)/2 left = (img_width - crop_size)/2 bottom = top + crop_size right = left + crop_size # Calculate the final coordinates for cropping crop_top = max(top, 0) crop_left = max(left, 0) crop_bottom = min(bottom, img_height) crop_right = min(right, img_width) # Ensure that the cropping region has non-zero width and height crop_width = crop_right - crop_left crop_height = crop_bottom - crop_top if crop_width <= 0 or crop_height <= 0: raise ValueError( "Invalid crop dimensions. Please check the values for top, left, right, and bottom.") # Crop the image and resize crop = image.crop((crop_left, crop_top, crop_right, crop_bottom)) crop = crop.resize( (((crop.size[0] // 8) * 8), ((crop.size[1] // 8) * 8))) return (pil2tensor(crop),) class OFFSEGSToImage: @classmethod def INPUT_TYPES(s): return {"required": { "segs": ("SEGS", ), }, "optional": { "fallback_image_opt": ("IMAGE", ), } } RETURN_TYPES = ("IMAGE",) FUNCTION = "doit" CATEGORY = "OFF" def doit(self, segs, fallback_image_opt=None): results = list() for seg in segs[1]: if seg.cropped_image is not None: cropped_image = torch.from_numpy(seg.cropped_image) elif fallback_image_opt is not None: # take from original image cropped_image = torch.from_numpy( crop_image(fallback_image_opt, seg.crop_region)) else: cropped_image = empty_pil_tensor() results.append(cropped_image) if len(results) == 0: results.append(empty_pil_tensor()) return (results[0],) class OFFCenterCropSEGS: @classmethod def INPUT_TYPES(s): return {"required": { "segs": ("SEGS", ), "image": ("IMAGE",) } } RETURN_TYPES = ("IMAGE",) FUNCTION = "doit" CATEGORY = "OFF" def doit(self, segs, image=None): results = list() image = tensor2pil(image) img_width, img_height = image.size crop_size = img_width if img_width > img_height: crop_size = img_height for seg in segs[1]: if seg.cropped_image is not None: center_x = (seg.crop_region[0] + seg.crop_region[2])/2 center_y = (seg.crop_region[1] + seg.crop_region[3])/2 crop_top = center_y - crop_size/2 crop_left = center_x - crop_size/2 crop_bottom = center_y + crop_size/2 crop_right = center_x + crop_size/2 if crop_top < 0: crop_bottom = crop_bottom - crop_top crop_top = 0 if crop_left < 0: crop_right = crop_right - crop_left crop_left = 0 if crop_bottom > img_height: crop_top = crop_top - (crop_bottom - img_height) crop_bottom = img_height if crop_right > img_width: crop_left = crop_left - (crop_right - img_width) crop_right = img_width cropped_image = image.crop( (crop_left, crop_top, crop_right, crop_bottom)) cropped_image = cropped_image.resize( (((cropped_image.size[0] // 8) * 8), ((cropped_image.size[1] // 8) * 8))) cropped_image = pil2tensor(cropped_image) else: cropped_image = empty_pil_tensor() results.append(cropped_image) if len(results) == 0: results.append(empty_pil_tensor()) return (results[0],) class OFFWatermark: def __init__(self): pass @classmethod def INPUT_TYPES(s): return { "required":{ "source":("IMAGE",), "watermark":("IMAGE",) } } RETURN_TYPES = ("IMAGE",) FUNCTION = "doit" CATEGORY = "OFF" def doit(self, source, watermark): source_image = tensor2pil(source) watermark_image = tensor2pil(watermark).convert() watermark_ratio = watermark_image.height/ watermark_image.width new_wm_with = source_image.width*0.15 watermark_image = watermark_image.resize((int(new_wm_with), int(new_wm_with*watermark_ratio))) source_image.paste(watermark_image,(10,10), watermark_image) return (pil2tensor(source_image),) class MaskToImageFallback: @classmethod def INPUT_TYPES(s): return { "required": { "mask": ("MASK",), "image": ("IMAGE",), } } CATEGORY = "OFF" RETURN_TYPES = ("IMAGE",) FUNCTION = "mask_to_image" def mask_to_image(self, mask, image): if mask is not None: result = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3) else: zero_tensor = torch.zeros_like(image) result = zero_tensor return (result,) def dilate_mask(mask, dilation_factor, iter=1): if dilation_factor == 0: return mask kernel = np.ones((abs(dilation_factor), abs(dilation_factor)), np.uint8) mask = cv2.UMat(mask) kernel = cv2.UMat(kernel) result = cv2.erode(mask, kernel, iter) return result.get() class MaskDilationForEachFace: @classmethod def INPUT_TYPES(s): return{ "required": { "mask": ("MASK",), "segs": ("SEGS",), "dilation_factor": ("FLOAT",{"min": 0.0, "max": 10.0, "step": 0.1, "default": 1.0}) } } CATEGORY = "OFF" RETURN_TYPES = ("MASK",) FUNCTION = "process" def process(self, mask, segs, dilation_factor): if mask.shape[0] is not len(segs[1]): return (mask,) mask = mask.clone() for i in range(0, mask.shape[0]): crop_mask = torch.zeros_like(mask[i]) crop_left = segs[1][i].crop_region[0] crop_right = segs[1][i].crop_region[2] crop_top = segs[1][i].crop_region[1] crop_bottom = segs[1][i].crop_region[3] crop_mask[:,segs[1][i].crop_region[1]:segs[1][i].crop_region[3] , segs[1][i].crop_region[0]:segs[1][i].crop_region[2]] = 255 pixel_count = (crop_right - crop_left) * (crop_bottom - crop_top) dilation = int(math.sqrt(pixel_count)/3.0/25.0 * dilation_factor) print(f"face {i} : pixel_count , dilation = ", pixel_count, dilation) mask[i] = mask[i] * crop_mask cv_mask = mask[i].squeeze(0).numpy() cv_mask = dilate_mask(cv_mask, dilation_factor=dilation) mask[i] = torch.from_numpy(cv_mask).unsqueeze(0) return (mask,) class SegsToFaceCropData: @classmethod def INPUT_TYPES(s): return{ "required":{ "segs":("SEGS",), "face_index":("INT", {"min":0, "default":0}), } } CATEGORY = "OFF" RETURN_TYPES = ("IMAGE","CROP_DATA",) FUNCTION ="process" def process(self, segs, face_index): return (segs[1][face_index].cropped_image, segs[1][face_index].crop_region,) class PasteFaceSegToImage: @classmethod def INPUT_TYPES(s): return { "required":{ "orig_image":("IMAGE",), "crop_image":("IMAGE",), "crop_data":("CROP_DATA",), } } CATEGORY = "OFF" RETURN_TYPES = ("IMAGE",) FUNCTION = "process" def process(self, orig_image, crop_image, crop_data): orig_image[:,crop_data[1]: crop_data[1]+crop_image.shape[1], crop_data[0]:crop_data[0]+crop_image.shape[2],:] = crop_image return (orig_image,) class OFFCLAHE: @classmethod def INPUT_TYPES(s): return{ "required":{ "image":("IMAGE",), "clip_limit":("FLOAT",{"default":1.2}), "grid_size":("INT",{"min":1, "default":8}) } } CATEGORY = "OFF" RETURN_TYPES = ("IMAGE",) FUNCTION = "process" def process(self, image, clip_limit, grid_size): batch_size, height, width, _ = image.shape result = torch.zeros_like(image) for i in range(batch_size): tensor_image = image[i].numpy() source_image = np.array(Image.fromarray((tensor_image*255).astype(np.uint8))) lab = cv2.cvtColor(source_image, cv2.COLOR_BGR2LAB) l, a, b = cv2.split(lab) clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=(grid_size, grid_size)) cl = clahe.apply(l) enhanced_lab = cv2.merge((cl, a, b)) enhanced_image = cv2.cvtColor(enhanced_lab, cv2.COLOR_LAB2BGR) result[i] = torch.from_numpy(enhanced_image.astype(np.float32) / 255.0) return (result,) class OffGridImageBatch: #from WAS SUITE def __init__(self): pass @classmethod def INPUT_TYPES(cls): return { "required": { "images": ("IMAGE",), "border_width": ("INT", {"default":20, "min": 0, "max": 100, "step":1}), "number_of_columns": ("INT", {"default":3, "min": 1, "max": 24, "step":1}), "max_cell_size": ("INT", {"default":320, "min":32, "max":2048, "step":1}), "border_red": ("INT", {"default":38, "min": 0, "max": 255, "step":1}), "border_green": ("INT", {"default":38, "min": 0, "max": 255, "step":1}), "border_blue": ("INT", {"default":38, "min": 0, "max": 255, "step":1}), "border_radius": ("INT", {"default":24, "min":0, "max":128,"step":1}), "target_size": ("INT", {"default":1080,"min":0,"max":8192,"step":8}) } } RETURN_TYPES = ("IMAGE",) FUNCTION = "smart_grid_image" CATEGORY = "OFF" def smart_grid_image(self, images, number_of_columns=3, max_cell_size=320, border_red=38, border_green=38, border_blue=38, border_width=20, border_radius=24, target_size=1080): cols = number_of_columns border_color = (border_red, border_green, border_blue) images_resized = [] max_row_height = 0 for tensor_img in images: img = tensor2pil(tensor_img) img_w, img_h = img.size aspect_ratio = img_w / img_h if img_w > img_h: cell_w = min(img_w, max_cell_size) cell_h = int(cell_w / aspect_ratio) else: cell_h = min(img_h, max_cell_size) cell_w = int(cell_h * aspect_ratio) img_resized = img.resize((cell_w, cell_h)) mask = Image.new("L", img_resized.size, 0) draw = ImageDraw.Draw(mask) draw.rounded_rectangle((0,0,cell_w, cell_h),fill=255,radius=border_radius) back_color = Image.new(img_resized.mode, img_resized.size, border_color) img_resized = Image.composite(img_resized, back_color, mask) img_resized = ImageOps.expand(img_resized, border=border_width // 2, fill=border_color) images_resized.append(img_resized) max_row_height = max(max_row_height, cell_h) max_row_height = int(max_row_height) total_images = len(images_resized) rows = math.ceil(total_images / cols) grid_width = cols * max_cell_size + (cols - 1) * border_width grid_height = rows * max_row_height + (rows - 1) * border_width new_image = Image.new('RGB', (grid_width, grid_height), border_color) rand_idx = np.arange(len(images_resized)) np.random.shuffle(rand_idx) for i, img in enumerate(images_resized): x = (rand_idx[i] % cols) * (max_cell_size + border_width) y = (rand_idx[i] // cols) * (max_row_height + border_width) img_w, img_h = img.size paste_x = x + (max_cell_size - img_w) // 2 paste_y = y + (max_row_height - img_h) // 2 new_image.paste(img, (paste_x, paste_y, paste_x + img_w, paste_y + img_h)) new_w, _ = new_image.size outer_border_width = int((target_size - new_w)/2) if outer_border_width>0 : new_image = ImageOps.expand(new_image, border=outer_border_width, fill=border_color) return (pil2tensor(new_image), ) class CalcMaskBound: def __init__(self): pass @classmethod def INPUT_TYPES(s): return { "required": { "mask": ("MASK",), "padding":("INT",{"default":0, "min": 0, "step":1}) } } CATEGORY = "OFF" RETURN_TYPES = ("INT","INT","INT","INT",) FUNCTION = "process" def process(self, mask, padding): rows, cols = np.where(mask[0]) min_row, max_row = rows.min(), rows.max() min_col, max_col = cols.min(), cols.max() min_col -= padding min_row -= padding max_col += padding max_row += padding if min_col< 0 : min_col = 0 if min_row<0: min_row = 0 if max_row > mask.shape[1]: max_row = mask.shape[1] if max_col > mask.shape[2]: max_col = mask.shape[2] return (min_col, min_row, max_col - min_col, max_row - min_row, )