nodes/EllangoK-ComfyUI-post-processing-nodes
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

slight refactoring of existing nodes

Browse Source
This commit is contained in:
EllangoK
2023-03-30 16:54:27 -04:00
parent 82d6271f4b
commit ece18e5f0d
6 changed files with 92 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
canny_edge_detect.py
+6 -6
View File
@@ -15,13 +15,13 @@ class CannyEdgeDetection:
"default": 100,
"min": 0,
"max": 500,
"step": 1
"step": 10
}),
"upper_threshold": ("INT", {
"default": 200,
"min": 0,
"max": 500,
"step": 1
"step": 10
}),
},
}
@@ -31,10 +31,10 @@ class CannyEdgeDetection:
CATEGORY = "postprocessing"
def canny(self, image, lower_threshold, upper_threshold):
tensor_img = image.numpy()[0]
tensor_img = (cv2.cvtColor(tensor_img, cv2.COLOR_BGR2GRAY) * 255).astype(np.uint8)
canny = cv2.Canny(tensor_img, lower_threshold, upper_threshold)
def canny(self, image: torch.Tensor, lower_threshold: int, upper_threshold: int):
tensor_image = image.numpy()[0]
gray_image = (cv2.cvtColor(tensor_image, cv2.COLOR_BGR2GRAY) * 255).astype(np.uint8)
canny = cv2.Canny(gray_image, lower_threshold, upper_threshold)
tensor = torch.from_numpy(canny).unsqueeze(0)
return (tensor,)
color_correct.py
+3 -6
View File
@@ -56,8 +56,8 @@ class ColorCorrect:
CATEGORY = "postprocessing"
def color_correct(self, image, temperature, hue, brightness, contrast, saturation, gamma):
tensor_img = image.numpy()[0]
def color_correct(self, image: torch.Tensor, temperature: float, hue: float, brightness: float, contrast: float, saturation: float, gamma: float):
tensor_image = image.numpy()[0]
brightness /= 100
contrast /= 100
@@ -68,8 +68,7 @@ class ColorCorrect:
contrast = 1 + contrast
saturation = 1 + saturation
modified_image = Image.fromarray((tensor_img * 255).astype(np.uint8))
modified_image = Image.fromarray((tensor_image * 255).astype(np.uint8))
# brightness
modified_image = ImageEnhance.Brightness(modified_image).enhance(brightness)
@@ -97,11 +96,9 @@ class ColorCorrect:
# hue
hsv_img = cv2.cvtColor(modified_image, cv2.COLOR_RGB2HSV)
hsv_img[:, :, 0] = (hsv_img[:, :, 0] + hue) % 360
modified_image = cv2.cvtColor(hsv_img, cv2.COLOR_HSV2RGB)
modified_image = modified_image.astype(np.uint8)
modified_image = modified_image / 255
modified_image = torch.from_numpy(modified_image).unsqueeze(0)
dither.py
+20 -22
View File
@@ -11,7 +11,7 @@ class Dither:
"image": ("IMAGE",),
"bits": ("INT", {
"default": 4,
"min": 0,
"min": 1,
"max": 8,
"step": 1
}),
@@ -23,36 +23,34 @@ class Dither:
CATEGORY = "postprocessing"
def dither(self, image, bits):
tensor_img = image[0]
height, width, _ = tensor_img.shape
out = tensor_img.clone()
levels = 2 ** bits - 1
def dither(self, image: torch.Tensor, bits: int):
tensor_image = image.numpy()[0]
img = (tensor_image * 255)
height, width, _ = img.shape
scale = 255 / (2**bits - 1)
for y in range(height):
for x in range(width):
old_pixel = out[y, x].clone()
new_pixel = torch.round(old_pixel * levels) / levels
out[y, x] = new_pixel
old_pixel = img[y, x].copy()
new_pixel = np.round(old_pixel / scale) * scale
img[y, x] = new_pixel
error = old_pixel - new_pixel
quant_error = old_pixel - new_pixel
if x + 1 < width:
out[y, x + 1] += error * (7 / 16)
if x - 1 >= 0 and y + 1 < height:
out[y + 1, x - 1] += error * 3/16
img[y, x + 1] += quant_error * 7 / 16
if y + 1 < height:
out[y + 1, x] += error * 5/16
if x + 1 < width and y + 1 < height:
out[y + 1, x + 1] += error * 1/16
if x - 1 >= 0:
img[y + 1, x - 1] += quant_error * 3 / 16
img[y + 1, x] += quant_error * 5 / 16
if x + 1 < width:
img[y + 1, x + 1] += quant_error * 1 / 16
out = torch.clamp(out, 0, 1).unsqueeze(0)
dithered = img / 255
tensor = torch.from_numpy(dithered).unsqueeze(0)
return (tensor,)
return (out,)
# A dictionary that contains all nodes you want to export with their names
# NOTE: names should be globally unique
NODE_CLASS_MAPPINGS = {
"Dither": Dither
}
gaussian_blur.py
+3 -3
View File
@@ -30,9 +30,9 @@ class GaussianBlur:
CATEGORY = "postprocessing"
def blur(self, image, kernel_size, sigma):
tensor_img = image.numpy()[0]
blurred = cv2.GaussianBlur(tensor_img, (kernel_size, kernel_size), sigma)
def blur(self, image: torch.Tensor, kernel_size: int, sigma: float):
tensor_image = image.numpy()[0]
blurred = cv2.GaussianBlur(tensor_image, (kernel_size, kernel_size), sigma)
tensor = torch.from_numpy(blurred).unsqueeze(0)
return (tensor,)
kmeans_quantize.py
+53
View File
@@ -0,0 +1,53 @@
import numpy as np
import cv2
import torch
class KMeansQuantize:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"colors": ("INT", {
"default": 16,
"min": 1,
"max": 256,
"step": 1
}),
"precision": ("INT", {
"default": 10,
"min": 1,
"max": 100,
"step": 1
}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "kmeans_quantize"
CATEGORY = "postprocessing"
def kmeans_quantize(self, image: torch.Tensor, colors: int, precision: int):
tensor_image = image.numpy()[0].astype(np.float32)
img = tensor_image
height, width, c = img.shape
criteria = (
cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER,
precision * 5, 0.01
)
img_copy = img.reshape(-1, c)
_, label, center = cv2.kmeans(
img_copy, colors, None,
criteria, 1, cv2.KMEANS_PP_CENTERS
)
result = center[label.flatten()].reshape(*img.shape)
tensor = torch.from_numpy(result).unsqueeze(0)
return (tensor,)
sharpen.py
+4 -3
View File
@@ -31,17 +31,18 @@ class Sharpen:
CATEGORY = "postprocessing"
def sharpen(self, image, alpha, kernel_size):
tensor_img = image.numpy()[0]
def sharpen(self, image: torch.Tensor, kernel_size: int, alpha: float):
tensor_image = image.numpy()[0]
kernel = np.ones((kernel_size, kernel_size), dtype=np.float32) * -1
center = kernel_size // 2
kernel[center, center] = kernel_size**2
kernel *= alpha
sharpened = cv2.filter2D(tensor_img, -1, kernel)
sharpened = cv2.filter2D(tensor_image, -1, kernel)
tensor = torch.from_numpy(sharpened).unsqueeze(0)
tensor = torch.clamp(tensor, 0, 1)
return (tensor,)
NODE_CLASS_MAPPINGS = {