60 lines
2.1 KiB
Python
60 lines
2.1 KiB
Python
import numpy as np
|
|
from .utils.color import extract_pad_color, get_empty_color, get_null_color
|
|
import torch
|
|
from .categories import CATE_TEST
|
|
|
|
|
|
class FEColor2Image:
|
|
"""
|
|
生成RGB颜色
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"colorLU": ("RGB_COLOR", {"default": get_empty_color()}),
|
|
"height": ("INT", {"default": 512, "min": 0, "max": 8192, "step": 8}),
|
|
"width": ("INT", {"default": 512, "min": 0, "max": 8192, "step": 8}),
|
|
"batch_size": ("INT", {"default": 1, "min": 1, "max": 64, "step": 1}),
|
|
"batch_rand": (['true', 'false'],)
|
|
},
|
|
"optional": {
|
|
"colorLD": ("RGB_COLOR", {"default": get_null_color()}),
|
|
"colorRU": ("RGB_COLOR", {"default": get_null_color()}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ('image',)
|
|
FUNCTION = "build_image"
|
|
CATEGORY = CATE_TEST
|
|
|
|
def build_image(self, colorLU, height, width, batch_size, colorLD=None, colorRU=None, batch_rand="false"):
|
|
if batch_rand == "true":
|
|
t_batch_size = batch_size
|
|
e_batch_size = -1
|
|
else:
|
|
t_batch_size = 1
|
|
e_batch_size = batch_size
|
|
color_lt = extract_pad_color(colorLU, t_batch_size)
|
|
color_lt = color_lt.expand(e_batch_size, 3)
|
|
if colorLD is None:
|
|
colorLD = colorLU
|
|
color_lb = extract_pad_color(colorLD, t_batch_size).expand(e_batch_size, 3)
|
|
if colorRU is None:
|
|
colorRU = colorLU
|
|
color_rt = extract_pad_color(colorRU, t_batch_size).expand(e_batch_size, 3)
|
|
|
|
image = torch.zeros((batch_size, height, width, 3), dtype=torch.float32)
|
|
|
|
weight_rt = torch.linspace(0, 1, steps=width).view(1, -1, 1)
|
|
weight_lb = torch.linspace(0, 1, steps=height).view(-1, 1, 1)
|
|
weight_lt = 1 - weight_rt - weight_lb
|
|
|
|
for _ in range(color_lt.shape[0]):
|
|
image[_, :, :, :] = weight_lt * color_lt[_] + weight_rt * color_rt[_] + weight_lb * color_lb[_]
|
|
|
|
image = image / 255.0
|
|
return (image,)
|