239 lines
8.6 KiB
Python
239 lines
8.6 KiB
Python
import torch
|
|
import random
|
|
import numpy as np
|
|
from enum import Enum
|
|
from nodes import MAX_RESOLUTION
|
|
|
|
# Images generated by combination of different creases need to be normalized
|
|
def normalize_img(img):
|
|
return (img / img.max()) * 255
|
|
|
|
# cross = True if cross crease is wanted
|
|
def create_uneven_array(low, up, steps, spacing=1, cross=False):
|
|
span = up - low
|
|
dx = 1.0 / steps
|
|
if cross:
|
|
arr = np.array([low + (i*dx)**spacing*span for i in range(steps//2)])
|
|
return np.append(arr, arr[::-1])
|
|
else :
|
|
arr = np.array([low + (i*dx)**spacing*span for i in range(steps)])
|
|
return arr
|
|
|
|
def parabolic_crease(spacing, c, scale=100, corner=1, resolution = 1000):
|
|
"""
|
|
Parameters:
|
|
|
|
spacing = controls how close the intermediate values will be to lower value
|
|
c = higher the c more spread out the gradient will be
|
|
scale = lesser the scale more concentrated is gradient towards the corner
|
|
"""
|
|
img = np.zeros((resolution, resolution))
|
|
|
|
# Varying the scaling parameter of create_uneven_array will give the parabolic gradient transition
|
|
for i in range(resolution):
|
|
img[i] = create_uneven_array(255, 0, resolution, spacing + c*i/scale)
|
|
|
|
if corner == 1:
|
|
return img
|
|
elif corner == 2:
|
|
return img[::-1]
|
|
elif corner == 3:
|
|
return img.T
|
|
else:
|
|
return img.T[::-1]
|
|
|
|
# If cross=1, then cross crease else linear crease is returned
|
|
def cross_crease(spacing, cross=1, resolution = 1000):
|
|
a = create_uneven_array(255, 0, resolution, spacing, cross=True)
|
|
img = np.tile(a, (resolution, 1))
|
|
return normalize_img(img*img.T) if cross else img
|
|
|
|
# Final function to return some random crease from 8 different types
|
|
def custom_crease():
|
|
spacing = random.uniform(1, 1.5)
|
|
scale = random.randint(100, 300)
|
|
corner = random.randint(1, 4)
|
|
|
|
# constant determines the type of crease and also is used to scale spacing in parabolic_crease
|
|
constant = random.randint(1, 10)
|
|
|
|
# Returning those creases which are based on parabolic
|
|
parabolic = parabolic_crease(spacing, constant, scale, corner)
|
|
if constant == 1:
|
|
return parabolic
|
|
elif constant == 2:
|
|
return normalize_img(parabolic*parabolic.T*parabolic[::-1]*parabolic.T[::-1])
|
|
|
|
# Returning those creases which are based on parabolic and cross
|
|
cross = cross_crease(spacing)
|
|
if constant == 3:
|
|
return cross
|
|
elif constant == 4:
|
|
return normalize_img(parabolic * cross)
|
|
elif constant == 5:
|
|
return normalize_img(cross * parabolic * parabolic.T)
|
|
|
|
# Returning those creases which are based on parabolic and linear
|
|
linear = cross_crease(spacing, 0)
|
|
if constant == 6:
|
|
return linear
|
|
elif constant == 7:
|
|
return linear.T
|
|
else:
|
|
return normalize_img(linear * parabolic)
|
|
|
|
|
|
|
|
class LightPosition(Enum):
|
|
LEFT = "Left Light"
|
|
RIGHT = "Right Light"
|
|
TOP = "Top Light"
|
|
BOTTOM = "Bottom Light"
|
|
TOP_LEFT = "Top Left Light"
|
|
TOP_RIGHT = "Top Right Light"
|
|
BOTTOM_LEFT = "Bottom Left Light"
|
|
BOTTOM_RIGHT = "Bottom Right Light"
|
|
|
|
def toRgb(color):
|
|
if color.startswith('#') and len(color) == 7: # e.g. "#RRGGBB"
|
|
color_rgb =tuple(int(color[i:i+2], 16) for i in (1, 3, 5))
|
|
else: # e.g. "255,255,255"
|
|
color_rgb = tuple(int(i) for i in color.split(','))
|
|
return color_rgb
|
|
|
|
def rgb_to_int(rgb):
|
|
r, g, b = rgb
|
|
return (r << 16) + (g << 8) + b
|
|
|
|
def generate_gradient_image(width:int, height:int, start_color:tuple=(255,255,255), end_color:tuple=(0,0,0), multiplier:float=1.0, lightPosition:LightPosition=LightPosition.LEFT):
|
|
"""
|
|
Generate a gradient image with a light source effect.
|
|
|
|
Parameters:
|
|
width (int): Width of the image.
|
|
height (int): Height of the image.
|
|
start_color: Starting color RGB of the gradient.
|
|
end_color: Ending color RGB of the gradient.
|
|
multiplier: Weight of light.
|
|
lightPosition (LightPosition): Position of the light source.
|
|
|
|
Returns:
|
|
np.array: 2D gradient image array.
|
|
"""
|
|
# Create a gradient from 0 to 1 and apply multiplier
|
|
if lightPosition == LightPosition.LEFT:
|
|
gradient = np.tile(np.linspace(0, 1, width)**multiplier, (height, 1))
|
|
elif lightPosition == LightPosition.RIGHT:
|
|
gradient = np.tile(np.linspace(1, 0, width)**multiplier, (height, 1))
|
|
elif lightPosition == LightPosition.TOP:
|
|
gradient = np.tile(np.linspace(0, 1, height)**multiplier, (width, 1)).T
|
|
elif lightPosition == LightPosition.BOTTOM:
|
|
gradient = np.tile(np.linspace(1, 0, height)**multiplier, (width, 1)).T
|
|
elif lightPosition == LightPosition.BOTTOM_RIGHT:
|
|
x = np.linspace(1, 0, width)**multiplier
|
|
y = np.linspace(1, 0, height)**multiplier
|
|
x_mesh, y_mesh = np.meshgrid(x, y)
|
|
gradient = np.sqrt(x_mesh**2 + y_mesh**2) / np.sqrt(2.0)
|
|
elif lightPosition == LightPosition.BOTTOM_LEFT:
|
|
x = np.linspace(0, 1, width)**multiplier
|
|
y = np.linspace(1, 0, height)**multiplier
|
|
x_mesh, y_mesh = np.meshgrid(x, y)
|
|
gradient = np.sqrt(x_mesh**2 + y_mesh**2) / np.sqrt(2.0)
|
|
elif lightPosition == LightPosition.TOP_RIGHT:
|
|
x = np.linspace(1, 0, width)**multiplier
|
|
y = np.linspace(0, 1, height)**multiplier
|
|
x_mesh, y_mesh = np.meshgrid(x, y)
|
|
gradient = np.sqrt(x_mesh**2 + y_mesh**2) / np.sqrt(2.0)
|
|
elif lightPosition == LightPosition.TOP_LEFT:
|
|
x = np.linspace(0, 1, width)**multiplier
|
|
y = np.linspace(0, 1, height)**multiplier
|
|
x_mesh, y_mesh = np.meshgrid(x, y)
|
|
gradient = np.sqrt(x_mesh**2 + y_mesh**2) / np.sqrt(2.0)
|
|
else:
|
|
raise ValueError(f"Unsupported position. Choose from {', '.join([member.value for member in LightPosition])}.")
|
|
|
|
# Interpolate between start_color and end_color based on the gradient
|
|
gradient_img = np.zeros((height, width, 3), dtype=np.float32)
|
|
for i in range(3):
|
|
gradient_img[..., i] = start_color[i] + (end_color[i] - start_color[i]) * gradient
|
|
|
|
gradient_img = np.clip(gradient_img, 0, 255).astype(np.uint8)
|
|
return gradient_img
|
|
|
|
|
|
def numpy_to_tensor(array: np.ndarray) -> torch.Tensor:
|
|
"""Convert a numpy array to a tensor and scale its values from 0-255 to 0-1."""
|
|
array = array.astype(np.float32) / 255.0
|
|
return torch.from_numpy(array)[None,]
|
|
|
|
|
|
class ImageGradient:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"light_position": ([member.value for member in LightPosition],),
|
|
"multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
|
"start_color": ("STRING", {"default": "#FFFFFF"}),
|
|
"end_color": ("STRING", {"default": "#000000"}),
|
|
"width": ("INT", { "default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 8, }),
|
|
"height": ("INT", { "default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 8, })
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "execute"
|
|
CATEGORY = "LightGradient"
|
|
DESCRIPTION = """Simple Light Gradient"""
|
|
|
|
def execute(self, light_position, multiplier, start_color, end_color, width, height):
|
|
lightPosition = LightPosition(light_position)
|
|
start_color_rgb = toRgb(start_color)
|
|
end_color_rgb = toRgb(end_color)
|
|
|
|
image = generate_gradient_image(width, height, start_color_rgb, end_color_rgb, multiplier, lightPosition)
|
|
image = numpy_to_tensor(image)
|
|
|
|
mask = generate_gradient_image(width, height, multiplier=multiplier, lightPosition=lightPosition)
|
|
mask = numpy_to_tensor(mask)
|
|
mask = mask[:, :, :, 0]
|
|
|
|
return (image,mask,)
|
|
|
|
class MaskGradient:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"light_position": ([member.value for member in LightPosition],),
|
|
"multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
|
"width": ("INT", { "default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 8, }),
|
|
"height": ("INT", { "default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 8, })
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MASK",)
|
|
FUNCTION = "execute"
|
|
CATEGORY = "LightGradient"
|
|
DESCRIPTION = """Mask Gradient"""
|
|
|
|
def execute(self, light_position, multiplier, width, height):
|
|
lightPosition = LightPosition(light_position)
|
|
|
|
mask = generate_gradient_image(width, height, multiplier=multiplier, lightPosition=lightPosition)
|
|
mask = numpy_to_tensor(mask)
|
|
mask = mask[:, :, :, 0]
|
|
|
|
return (mask,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"ImageGradient": ImageGradient,
|
|
"MaskGradient": MaskGradient
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"ImageGradient": "Image Gradient",
|
|
"MaskGradient": "Mask Gradient"
|
|
}
|