File size: 2,304 Bytes
aa289db |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
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",), # 输入的tensor图像
"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:
# Convert tensor to PIL for processing
_canvas = tensor2pil(torch.unsqueeze(i, 0)).convert('RGB')
# Convert PIL to OpenCV format
img_np = np.array(_canvas)[:, :, ::-1]
# 获取原始尺寸
height, width = img_np.shape[:2]
# 计算新尺寸
new_width = int(width / scale_factor)
new_height = int(height / scale_factor)
# 1. 使用cv2.INTER_AREA方法缩小图像
resized_down = cv2.resize(img_np, (new_width, new_height), interpolation=cv2.INTER_AREA)
# 2. 使用linear方法放大回原尺寸
resized_img = cv2.resize(resized_down, (width, height), interpolation=cv2.INTER_CUBIC)
# Convert OpenCV back to PIL and then to tensor
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"
}
|