|
import cv2 |
|
import numpy as np |
|
import torch |
|
from PIL import Image |
|
|
|
NODE_NAME = 'TTPlanet_Tile_Preprocessor' |
|
|
|
|
|
def pil2tensor(image: Image) -> torch.Tensor: |
|
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) |
|
|
|
def tensor2pil(t_image: torch.Tensor) -> Image: |
|
return Image.fromarray(np.clip(255.0 * t_image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)) |
|
|
|
class TTPlanet_Tile_Preprocessor: |
|
def __init__(self): |
|
pass |
|
|
|
@classmethod |
|
def INPUT_TYPES(cls): |
|
return { |
|
"required": { |
|
"image": ("IMAGE",), |
|
"scale_factor": ("FLOAT", {"default": 2.0, "min": 1.0, "max": 8.0, "step": 0.1}), |
|
}, |
|
"optional": {} |
|
} |
|
|
|
RETURN_TYPES = ("IMAGE",) |
|
RETURN_NAMES = ("image_output",) |
|
FUNCTION = 'process_image' |
|
CATEGORY = 'TTP_TILE' |
|
|
|
def process_image(self, image, scale_factor): |
|
ret_images = [] |
|
|
|
for i in image: |
|
|
|
_canvas = tensor2pil(torch.unsqueeze(i, 0)).convert('RGB') |
|
|
|
|
|
img_np = np.array(_canvas)[:, :, ::-1] |
|
|
|
|
|
height, width = img_np.shape[:2] |
|
|
|
|
|
new_width = int(width / scale_factor) |
|
new_height = int(height / scale_factor) |
|
|
|
|
|
resized_down = cv2.resize(img_np, (new_width, new_height), interpolation=cv2.INTER_AREA) |
|
|
|
|
|
resized_img = cv2.resize(resized_down, (width, height), interpolation=cv2.INTER_CUBIC) |
|
|
|
|
|
pil_img = Image.fromarray(resized_img[:, :, ::-1]) |
|
tensor_img = pil2tensor(pil_img) |
|
ret_images.append(tensor_img) |
|
|
|
return (torch.cat(ret_images, dim=0),) |
|
|
|
NODE_CLASS_MAPPINGS = { |
|
"Image Processing: TTPlanet_Tile_Preprocessor": TTPlanet_Tile_Preprocessor |
|
} |
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = { |
|
"Image Processing: TTPlanet_Tile_Preprocessor": "TTPlanet Tile Preprocessor" |
|
} |
|
|