Fix aspect ration on voronoi mask and add initial version of magic mask

This commit is contained in:
kijai 2023-11-20 00:03:24 +02:00
parent 02b00c6480
commit 3569b842df
2 changed files with 188 additions and 5 deletions

95
magictex.py Normal file
View File

@ -0,0 +1,95 @@
"""Generates psychedelic color textures in the spirit of Blender's magic texture shader using Python/Numpy
https://github.com/cheind/magic-texture
"""
from typing import Tuple, Optional
import numpy as np
def coordinate_grid(shape: Tuple[int, int], dtype=np.float32):
"""Returns a three-dimensional coordinate grid of given shape for use in `magic`."""
x = np.linspace(-1, 1, shape[1], endpoint=True, dtype=dtype)
y = np.linspace(-1, 1, shape[0], endpoint=True, dtype=dtype)
X, Y = np.meshgrid(x, y)
XYZ = np.stack((X, Y, np.ones_like(X)), -1)
return XYZ
def random_transform(coords: np.ndarray, rng: np.random.Generator = None):
"""Returns randomly transformed coordinates"""
H, W = coords.shape[:2]
rng = rng or np.random.default_rng()
m = rng.uniform(-1.0, 1.0, size=(3, 3)).astype(coords.dtype)
return (coords.reshape(-1, 3) @ m.T).reshape(H, W, 3)
def magic(
coords: np.ndarray,
depth: Optional[int] = None,
distortion: Optional[int] = None,
rng: np.random.Generator = None,
):
"""Returns color magic color texture.
The implementation is based on Blender's (https://www.blender.org/) magic
texture shader. The following adaptions have been made:
- we exchange the nested if-cascade by a probabilistic iterative approach
Kwargs
------
coords: HxWx3 array
Coordinates transformed into colors by this method. See
`magictex.coordinate_grid` to generate the default.
depth: int (optional)
Number of transformations applied. Higher numbers lead to more
nested patterns. If not specified, randomly sampled.
distortion: float (optional)
Distortion of patterns. Larger values indicate more distortion,
lower values tend to generate smoother patterns. If not specified,
randomly sampled.
rng: np.random.Generator
Optional random generator to draw samples from.
Returns
-------
colors: HxWx3 array
Three channel color image in range [0,1]
"""
rng = rng or np.random.default_rng()
if distortion is None:
distortion = rng.uniform(1, 4)
if depth is None:
depth = rng.integers(1, 5)
H, W = coords.shape[:2]
XYZ = coords
x = np.sin((XYZ[..., 0] + XYZ[..., 1] + XYZ[..., 2]) * distortion)
y = np.cos((-XYZ[..., 0] + XYZ[..., 1] - XYZ[..., 2]) * distortion)
z = -np.cos((-XYZ[..., 0] - XYZ[..., 1] + XYZ[..., 2]) * distortion)
if depth > 0:
x *= distortion
y *= distortion
z *= distortion
y = -np.cos(x - y + z)
y *= distortion
xyz = [x, y, z]
fns = [np.cos, np.sin]
for _ in range(1, depth):
axis = rng.choice(3)
fn = fns[rng.choice(2)]
signs = rng.binomial(n=1, p=0.5, size=4) * 2 - 1
xyz[axis] = signs[-1] * fn(
signs[0] * xyz[0] + signs[1] * xyz[1] + signs[2] * xyz[2]
)
xyz[axis] *= distortion
x, y, z = xyz
x /= 2 * distortion
y /= 2 * distortion
z /= 2 * distortion
c = 0.5 - np.stack((x, y, z), -1)
np.clip(c, 0, 1.0)
return c

View File

@ -1957,7 +1957,7 @@ class OffsetMask:
batch_size, height, width = mask.shape
if angle is not 0 and incremental:
if angle != 0 and incremental:
for i in range(batch_size):
rotation_angle = angle * (i+1)
mask[i] = TF.rotate(mask[i].unsqueeze(0), rotation_angle).squeeze(0)
@ -2128,7 +2128,7 @@ class CreateVoronoiMask:
def INPUT_TYPES(s):
return {
"required": {
"frames": ("INT", {"default": 1,"min": 1, "max": 4096, "step": 1}),
"frames": ("INT", {"default": 16,"min": 2, "max": 4096, "step": 1}),
"num_points": ("INT", {"default": 15,"min": 1, "max": 4096, "step": 1}),
"line_width": ("INT", {"default": 4,"min": 1, "max": 4096, "step": 1}),
"speed": ("FLOAT", {"default": 0.5,"min": 0.0, "max": 1.0, "step": 0.01}),
@ -2142,9 +2142,15 @@ class CreateVoronoiMask:
batch_size = frames
out = []
# Create start and end points for each point
# Calculate aspect ratio
aspect_ratio = frame_width / frame_height
# Create start and end points for each point, considering the aspect ratio
start_points = np.random.rand(num_points, 2)
start_points[:, 0] *= aspect_ratio
end_points = np.random.rand(num_points, 2)
end_points[:, 0] *= aspect_ratio
for i in range(batch_size):
# Interpolate the points' positions based on the current frame
@ -2152,15 +2158,18 @@ class CreateVoronoiMask:
t = np.clip(t, 0, 1) # ensure t is in [0, 1]
points = (1 - t) * start_points + t * end_points # lerp
# Adjust points for aspect ratio
points[:, 0] *= aspect_ratio
vor = Voronoi(points)
# Create a blank image with a white background
fig, ax = plt.subplots()
plt.subplots_adjust(left=0, right=1, bottom=0, top=1)
ax.set_xlim([0, 1]); ax.set_ylim([0, 1])
ax.set_xlim([0, aspect_ratio]); ax.set_ylim([0, 1]) # adjust x limits
ax.axis('off')
ax.margins(0, 0)
fig.set_size_inches(frame_width/100, frame_height/100)
fig.set_size_inches(aspect_ratio * frame_height/100, frame_height/100) # adjust figure size
ax.fill_between([0, 1], [0, 1], color='white')
# Plot each Voronoi ridge
@ -2181,6 +2190,83 @@ class CreateVoronoiMask:
return (torch.stack(out, dim=0), 1.0 - torch.stack(out, dim=0),)
from mpl_toolkits.axes_grid1 import ImageGrid
from .magictex import *
class CreateMagicMask:
RETURN_TYPES = ("MASK", "MASK",)
RETURN_NAMES = ("mask", "mask_inverted",)
FUNCTION = "createmagicmask"
CATEGORY = "KJNodes/masking/generate"
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"frames": ("INT", {"default": 16,"min": 2, "max": 4096, "step": 1}),
"depth": ("INT", {"default": 2,"min": 1, "max": 50, "step": 1}),
"distortion": ("FLOAT", {"default": 1.0,"min": 0.0, "max": 10.0, "step": 0.01}),
"seed": ("INT", {"default": 123,"min": 0, "max": 99999999, "step": 1}),
"transitions": ("INT", {"default": 2,"min": 1, "max": 20, "step": 1}),
"frame_width": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}),
"frame_height": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}),
},
}
def createmagicmask(self, frames, transitions, depth, distortion, seed, frame_width, frame_height):
rng = np.random.default_rng(seed)
out = []
coords = coordinate_grid((frame_width, frame_height))
# Calculate the number of frames for each transition
frames_per_transition = frames // transitions
# Generate a base set of parameters
base_params = {
"coords": random_transform(coords, rng),
"depth": depth,
"distortion": distortion,
}
for t in range(transitions):
# Generate a second set of parameters that is at most max_diff away from the base parameters
params1 = base_params.copy()
params2 = base_params.copy()
params1['coords'] = random_transform(coords, rng)
params2['coords'] = random_transform(coords, rng)
for i in range(frames_per_transition):
# Compute the interpolation factor
alpha = i / frames_per_transition
# Interpolate between the two sets of parameters
params = params1.copy()
params['coords'] = (1 - alpha) * params1['coords'] + alpha * params2['coords']
tex = magic(**params)
fig = plt.figure(figsize=(10, 10))
ax = fig.add_subplot(111)
plt.subplots_adjust(left=0, right=1, bottom=0, top=1)
ax.get_yaxis().set_ticks([])
ax.get_xaxis().set_ticks([])
ax.imshow(tex, aspect='auto')
fig.canvas.draw()
img = np.array(fig.canvas.renderer._renderer)
plt.close(fig)
pil_img = Image.fromarray(img).convert("L")
mask = torch.tensor(np.array(pil_img)) / 255.0
out.append(mask)
return (torch.stack(out, dim=0), 1.0 - torch.stack(out, dim=0),)
NODE_CLASS_MAPPINGS = {
"INTConstant": INTConstant,
"FloatConstant": FloatConstant,
@ -2220,6 +2306,7 @@ NODE_CLASS_MAPPINGS = {
"WidgetToString": WidgetToString,
"CreateShapeMask": CreateShapeMask,
"CreateVoronoiMask": CreateVoronoiMask,
"CreateMagicMask": CreateMagicMask,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"INTConstant": "INT Constant",
@ -2259,4 +2346,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"WidgetToString": "WidgetToString",
"CreateShapeMask": "CreateShapeMask",
"CreateVoronoiMask": "CreateVoronoiMask",
"CreateMagicMask": "CreateMagicMask",
}