297 lines
8.4 KiB
Python
297 lines
8.4 KiB
Python
import numpy as np
|
|
import torch
|
|
|
|
|
|
def channels_layer(images, channels, function):
|
|
if channels == "rgb":
|
|
img = images[:, :, :, :3]
|
|
elif channels == "rgba":
|
|
img = images[:, :, :, :4]
|
|
elif channels == "rg":
|
|
img = images[:, :, :, [0, 1]]
|
|
elif channels == "rb":
|
|
img = images[:, :, :, [0, 2]]
|
|
elif channels == "ra":
|
|
img = images[:, :, :, [0, 3]]
|
|
elif channels == "gb":
|
|
img = images[:, :, :, [1, 2]]
|
|
elif channels == "ga":
|
|
img = images[:, :, :, [1, 3]]
|
|
elif channels == "ba":
|
|
img = images[:, :, :, [2, 3]]
|
|
elif channels == "r":
|
|
img = images[:, :, :, 0]
|
|
elif channels == "g":
|
|
img = images[:, :, :, 1]
|
|
elif channels == "b":
|
|
img = images[:, :, :, 2]
|
|
elif channels == "a":
|
|
img = images[:, :, :, 3]
|
|
else:
|
|
raise ValueError("Unsupported channels")
|
|
|
|
result = torch.from_numpy(function(img.numpy()))
|
|
|
|
if channels == "rgb":
|
|
images[:, :, :, :3] = result
|
|
elif channels == "rgba":
|
|
images[:, :, :, :4] = result
|
|
elif channels == "rg":
|
|
images[:, :, :, [0, 1]] = result
|
|
elif channels == "rb":
|
|
images[:, :, :, [0, 2]] = result
|
|
elif channels == "ra":
|
|
images[:, :, :, [0, 3]] = result
|
|
elif channels == "gb":
|
|
images[:, :, :, [1, 2]] = result
|
|
elif channels == "ga":
|
|
images[:, :, :, [1, 3]] = result
|
|
elif channels == "ba":
|
|
images[:, :, :, [2, 3]] = result
|
|
elif channels == "r":
|
|
images[:, :, :, 0] = result
|
|
elif channels == "g":
|
|
images[:, :, :, 1] = result
|
|
elif channels == "b":
|
|
images[:, :, :, 2] = result
|
|
elif channels == "a":
|
|
images[:, :, :, 3] = result
|
|
|
|
return images
|
|
|
|
|
|
class ImageNoiseBeta:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"a": ("INT", {
|
|
"default": 1,
|
|
"min": 1,
|
|
}),
|
|
"b": ("INT", {
|
|
"default": 1,
|
|
"min": 1,
|
|
}),
|
|
"monochromatic": (["false", "true"],),
|
|
"invert": (["false", "true"],),
|
|
"channels": (["rgb", "rgba", "rg", "rb", "ra", "gb", "ga", "ba", "r", "g", "b", "a"],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "node"
|
|
CATEGORY = "image/noise"
|
|
|
|
def noise(self, images, a, b, monochromatic, invert):
|
|
if monochromatic and images.shape[3] > 1:
|
|
noise = np.random.beta(a, b, images.shape[:3])
|
|
else:
|
|
noise = np.random.beta(a, b, images.shape)
|
|
|
|
if monochromatic and images.shape[3] > 1:
|
|
noise = noise[..., np.newaxis].repeat(images.shape[3], -1)
|
|
|
|
if invert:
|
|
noise = images - noise
|
|
else:
|
|
noise = images + noise
|
|
|
|
noise = noise.astype(images.dtype)
|
|
|
|
return noise
|
|
|
|
def node(self, images, a, b, monochromatic, invert, channels):
|
|
tensor = images.clone().detach()
|
|
|
|
monochromatic = True if monochromatic == "true" else False
|
|
invert = True if invert == "true" else False
|
|
|
|
return (channels_layer(tensor, channels, lambda x: self.noise(
|
|
x, a, b, monochromatic, invert
|
|
)),)
|
|
|
|
|
|
class ImageNoiseBinomial:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"n": ("INT", {
|
|
"default": 128,
|
|
"min": 1,
|
|
"max": 255,
|
|
}),
|
|
"p": ("FLOAT", {
|
|
"default": 0.5,
|
|
"max": 1.0,
|
|
"step": 0.01
|
|
}),
|
|
"monochromatic": (["false", "true"],),
|
|
"invert": (["false", "true"],),
|
|
"channels": (["rgb", "rgba", "rg", "rb", "ra", "gb", "ga", "ba", "r", "g", "b", "a"],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "node"
|
|
CATEGORY = "image/noise"
|
|
|
|
def noise(self, images, n, p, monochromatic, invert):
|
|
if monochromatic and images.shape[3] > 1:
|
|
noise = np.random.binomial(n, p, images.shape[:3])
|
|
else:
|
|
noise = np.random.binomial(n, p, images.shape)
|
|
|
|
noise = noise.astype(images.dtype)
|
|
noise /= 255
|
|
|
|
if monochromatic and images.shape[3] > 1:
|
|
noise = noise[..., np.newaxis].repeat(images.shape[3], -1)
|
|
|
|
if invert:
|
|
noise = images - noise
|
|
else:
|
|
noise = images + noise
|
|
|
|
noise = np.clip(noise, 0.0, 1.0)
|
|
|
|
return noise
|
|
|
|
def node(self, images, n, p, monochromatic, invert, channels):
|
|
tensor = images.clone().detach()
|
|
|
|
monochromatic = True if monochromatic == "true" else False
|
|
invert = True if invert == "true" else False
|
|
|
|
return (channels_layer(tensor, channels, lambda x: self.noise(
|
|
x, n, p, monochromatic, invert
|
|
)),)
|
|
|
|
|
|
class ImageNoiseBytes:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"monochromatic": (["false", "true"],),
|
|
"invert": (["false", "true"],),
|
|
"channels": (["rgb", "rgba", "rg", "rb", "ra", "gb", "ga", "ba", "r", "g", "b", "a"],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "node"
|
|
CATEGORY = "image/noise"
|
|
|
|
def noise(self, images, monochromatic, invert):
|
|
if monochromatic and images.shape[3] > 1:
|
|
noise = np.random.bytes(np.prod(images.shape[:3]))
|
|
noise = np.frombuffer(noise, np.uint8)
|
|
noise = np.reshape(noise, images.shape[:3])
|
|
else:
|
|
noise = np.random.bytes(np.prod(images.shape))
|
|
noise = np.frombuffer(noise, np.uint8)
|
|
noise = np.reshape(noise, images.shape)
|
|
|
|
noise = noise.astype(images.dtype)
|
|
noise /= 255
|
|
|
|
if monochromatic and images.shape[3] > 1:
|
|
noise = noise[..., np.newaxis].repeat(images.shape[3], -1)
|
|
|
|
if invert:
|
|
noise = images - noise
|
|
else:
|
|
noise = images + noise
|
|
|
|
noise = np.clip(noise, 0.0, 1.0)
|
|
|
|
return noise
|
|
|
|
def node(self, images, monochromatic, invert, channels):
|
|
tensor = images.clone().detach()
|
|
|
|
monochromatic = True if monochromatic == "true" else False
|
|
invert = True if invert == "true" else False
|
|
|
|
return (channels_layer(tensor, channels, lambda x: self.noise(
|
|
x, monochromatic, invert
|
|
)),)
|
|
|
|
|
|
class ImageNoiseGaussian:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"strength": ("FLOAT", {
|
|
"default": 0.5,
|
|
"step": 0.01
|
|
}),
|
|
"monochromatic": (["false", "true"],),
|
|
"invert": (["false", "true"],),
|
|
"channels": (["rgb", "rgba", "rg", "rb", "ra", "gb", "ga", "ba", "r", "g", "b", "a"],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "node"
|
|
CATEGORY = "image/noise"
|
|
|
|
def noise(self, images, strength, monochromatic, invert):
|
|
if monochromatic and images.shape[3] > 1:
|
|
noise = np.random.normal(0, 1, images.shape[:3])
|
|
else:
|
|
noise = np.random.normal(0, 1, images.shape)
|
|
|
|
noise = np.abs(noise)
|
|
noise /= noise.max()
|
|
|
|
if monochromatic and images.shape[3] > 1:
|
|
noise = noise[..., np.newaxis].repeat(images.shape[3], -1)
|
|
|
|
if invert:
|
|
noise = images - noise * strength
|
|
else:
|
|
noise = images + noise * strength
|
|
|
|
noise = np.clip(noise, 0.0, 1.0)
|
|
noise = noise.astype(images.dtype)
|
|
|
|
return noise
|
|
|
|
def node(self, images, strength, monochromatic, invert, channels):
|
|
tensor = images.clone().detach()
|
|
|
|
monochromatic = True if monochromatic == "true" else False
|
|
invert = True if invert == "true" else False
|
|
|
|
return (channels_layer(tensor, channels, lambda x: self.noise(
|
|
x, strength, monochromatic, invert
|
|
)),)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"ImageNoiseBeta": ImageNoiseBeta,
|
|
"ImageNoiseBinomial": ImageNoiseBinomial,
|
|
"ImageNoiseBytes": ImageNoiseBytes,
|
|
"ImageNoiseGaussian": ImageNoiseGaussian
|
|
}
|