From 6e77293a3c8ece2e77ec37e46bb1e2f8ce73aea5 Mon Sep 17 00:00:00 2001 From: aravind Date: Mon, 18 Mar 2024 16:03:38 +0530 Subject: [PATCH] Merged scaled paste --- __init__.py | 9 ++- scaled_paste.py | 176 ++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 182 insertions(+), 3 deletions(-) create mode 100644 scaled_paste.py diff --git a/__init__.py b/__init__.py index 2cb8e38..89bc8e1 100644 --- a/__init__.py +++ b/__init__.py @@ -11,6 +11,7 @@ import hashlib import comfy.model_management as model_management import folder_paths from PIL import Image, ImageOps +from scaled_paste import main_scaled_paste def from_torch_image(image): @@ -2321,7 +2322,7 @@ class TRI3D_image_mask_2_box: FUNCTION = "run" RETURN_TYPES = ("IMAGE", ) - CATEGORY = "HackNode" + CATEGORY = "TRI3D" def run(self, image, mask): image = from_torch_image(image) @@ -2348,7 +2349,7 @@ class TRI3D_image_mask_box_2_image: FUNCTION = "run" RETURN_TYPES = ("IMAGE", ) - CATEGORY = "HackNode" + CATEGORY = "TRI3D" def run(self, image, mask, box): image = from_torch_image(image) @@ -2601,7 +2602,7 @@ class main_transparent_background(): FUNCTION = "run" RETURN_TYPES = ("IMAGE", "MASK") - CATEGORY = "HackNode" + CATEGORY = "TRI3D" def run(self, image): image = self.from_torch_image(image) @@ -2644,6 +2645,7 @@ NODE_CLASS_MAPPINGS = { "tri3d-HistogramEqualization": HistogramEqualization, "tri3d-composite-image-splitter": TRI3DCompositeImageSplitter, 'tri3d-main_transparent_background': main_transparent_background, + 'tri3d-scaled-paste': main_scaled_paste, } VERSION = "2.9.0" @@ -2675,4 +2677,5 @@ NODE_DISPLAY_NAME_MAPPINGS = { "tri3d-HistogramEqualization": "Adjust Neck" + " v" + VERSION, "tri3d-composite-image-splitter": "Composite Image Splitter" + " v" + VERSION, 'tri3d-main_transparent_background': 'Transparent Background' + " v" + VERSION, + 'tri3d-scaled-paste': 'Scaled paste' + " v" + VERSION, } diff --git a/scaled_paste.py b/scaled_paste.py new file mode 100644 index 0000000..833b207 --- /dev/null +++ b/scaled_paste.py @@ -0,0 +1,176 @@ +#!/usr/bin/python3 +import torch +import cv2 +import numpy as np + + +#!/usr/bin/python3 +def from_torch_image(image): + image = image.cpu().numpy() * 255.0 + image = np.clip(image, 0, 255).astype(np.uint8) + return image + + +def to_torch_image(image): + image = image.astype(dtype=np.float32) + image /= 255.0 + image = torch.from_numpy(image) + return image + + +def scaled_paste( + image_background, + image_foreground, + mask_foreground, + scale_factor, + height_factor=1.2, +): + + print('DEBUG scaled_paste 0 ', image_background.shape, + image_foreground.shape, mask_foreground.shape, scale_factor, + height_factor) + + height = image_foreground.shape[0] * height_factor + + print('DEBUG scaled_paste 1 ', height) + + max_0 = max(image_background.shape[0], height) + max_1 = max(image_background.shape[1], image_foreground.shape[1]) + + print('DEBUG scaled_paste 2 ', max_0, max_1) + + ratio_0 = max_0 / image_background.shape[0] + ratio_1 = max_1 / image_background.shape[1] + ratio_max = max(ratio_0, ratio_1) * scale_factor + + print('DEBUG scaled_paste 2 ', ratio_0, ratio_1, ratio_max) + + size_0 = int(image_background.shape[0] * ratio_max) + 1 + size_1 = int(image_background.shape[1] * ratio_max) + 1 + + print('DEBUG scaled_paste 3 ', size_0, size_1) + + image_background = cv2.resize(image_background, (size_1, size_0), + cv2.INTER_CUBIC) + + print('DEBUG scaled_paste 4 ', image_background.shape) + + end_0 = int(image_background.shape[0]) + begin_0 = int(end_0 - height) + end_0 = int(begin_0 + image_foreground.shape[0]) + + print('DEBUG scaled_paste 5 ', begin_0, end_0) + + end_1 = image_background.shape[1] + begin_1 = end_1 - image_foreground.shape[1] + begin_1 = int(begin_1 / 2) + end_1 = int(begin_1 + image_foreground.shape[1]) + + print('DEBUG scaled_paste 6 ', begin_1, end_1) + + image_reference = image_background[begin_0:end_0, begin_1:end_1, :] + + for i in range(3): + image_reference[:, :, + i] = (mask_foreground * image_foreground[:, :, i]) + ( + (1 - mask_foreground) * image_reference[:, :, i]) + + return image_background + + +#!/usr/bin/python3 +class main_scaled_paste(): + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image_background": ("IMAGE", ), + "image_foreground": ("IMAGE", ), + "mask_foreground": ("MASK", ), + "scale_factor": ("FLOAT", { + "default": 1.2, + "min": 1, + "max": 10, + "step": 0.05 + }), + "height_factor": ("FLOAT", { + "default": 1.01, + "min": 1, + "max": 8, + "step": 0.05 + }), + }, + } + + FUNCTION = "run" + RETURN_TYPES = ("IMAGE", ) + CATEGORY = "TRI3D" + + def run( + self, + image_background, + image_foreground, + mask_foreground, + scale_factor, + height_factor, + ): + + print('DEBUG 0 ', image_background.shape, image_foreground.shape, + mask_foreground.shape) + + image_background = from_torch_image(image_background) + image_foreground = from_torch_image(image_foreground) + mask_foreground = mask_foreground.cpu().numpy() + + image_output = scaled_paste( + image_background[0], + image_foreground[0], + mask_foreground[0], + scale_factor, + height_factor, + ) + + print('DEBUG 1 ', image_output.shape) + + image_output = to_torch_image(image=image_output) + + print('DEBUG 2 ', image_output.shape) + + image_output = image_output.unsqueeze(0) + + print('DEBUG 3 ', image_output.shape) + + return (image_output, ) + + +#!/usr/bin/python3 +# mask = cv2.imread('/home/asd/DATASETS/BG_SWAP_HACK_TEST/FOREGROUND_MASK.png', +# cv2.IMREAD_GRAYSCALE) + +# mask = mask.astype(dtype=np.float32) / 255.0 + +# image_background = scaled_paste( +# image_background=cv2.imread( +# '/home/asd/DATASETS/BG_SWAP_HACK_TEST/BACKGROUND_DEPTH.png', +# cv2.IMREAD_COLOR), +# image_foreground=cv2.imread( +# '/home/asd/DATASETS/BG_SWAP_HACK_TEST/FOREGROUND_DEPTH.png', +# cv2.IMREAD_COLOR), +# mask_foreground=mask, +# scale_factor=2, +# height_factor=1.05, +# ) + +# cv2.imwrite('tmp.png', image_background) + +NODE_CLASS_MAPPINGS = { + 'main_scaled_paste': main_scaled_paste, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + 'main_scaled_paste': 'main_scaled_paste', +}