67 lines
2.6 KiB
Python
67 lines
2.6 KiB
Python
from PIL import Image, ImageDraw
|
|
import numpy as np
|
|
import torch
|
|
import math
|
|
|
|
class OpticalIllusionNode:
|
|
CATEGORY = "illusion"
|
|
FUNCTION = "generate_illusion"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"illusion_type": (["checkerboard", "circles", "lines", "spiral"], {"default": "checkerboard"}),
|
|
"size": ("INT", {"default": 512, "min": 128, "max": 2048}),
|
|
"frequency": ("INT", {"default": 10, "min": 2, "max": 100}),
|
|
"line_width": ("INT", {"default": 3, "min": 1, "max": 100}),
|
|
"color1": ("STRING", {"default": "#FFFFFF"}),
|
|
"color2": ("STRING", {"default": "#000000"})
|
|
}
|
|
}
|
|
|
|
def generate_illusion(self, illusion_type, size, frequency, line_width, color1, color2):
|
|
img = Image.new('RGB', (size, size), color1)
|
|
draw = ImageDraw.Draw(img)
|
|
|
|
if illusion_type == "checkerboard":
|
|
tile = size // frequency
|
|
for y in range(frequency):
|
|
for x in range(frequency):
|
|
if (x + y) % 2 == 0:
|
|
draw.rectangle([x*tile, y*tile, (x+1)*tile, (y+1)*tile], fill=color2)
|
|
|
|
elif illusion_type == "circles":
|
|
step = size / (frequency * 2)
|
|
for i in range(frequency):
|
|
radius = step * (i + 1)
|
|
bbox = [size//2 - radius, size//2 - radius, size//2 + radius, size//2 + radius]
|
|
draw.ellipse(bbox, outline=color2 if i % 2 == 0 else color1, width=line_width)
|
|
|
|
elif illusion_type == "lines":
|
|
spacing = size / frequency
|
|
for i in range(frequency):
|
|
offset = i * spacing
|
|
draw.line([(offset, 0), (offset, size)], fill=color2 if i % 2 == 0 else color1, width=line_width)
|
|
|
|
elif illusion_type == "spiral":
|
|
cx, cy = size // 2, size // 2
|
|
max_radius = size * 0.48
|
|
num_turns = frequency
|
|
step_theta = math.pi / 720 # très fin = très lisse
|
|
a = 0
|
|
b = max_radius / (2 * math.pi * num_turns)
|
|
theta = 0
|
|
while theta < 2 * math.pi * num_turns:
|
|
r = a + b * theta
|
|
bbox = [cx - r, cy - r, cx + r, cy + r]
|
|
start = math.degrees(theta)
|
|
end = math.degrees(theta + step_theta)
|
|
draw.arc(bbox, start, end, fill=color2, width=line_width)
|
|
theta += step_theta
|
|
|
|
img_array = np.array(img).astype(np.float32) / 255.0
|
|
tensor = torch.from_numpy(img_array).unsqueeze(0)
|
|
return (tensor,)
|