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

allow for looping over the batch dim

Browse Source
This commit is contained in:
EllangoK
2023-03-30 19:51:10 -04:00
parent ece18e5f0d
commit 228c85c450
8 changed files with 153 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+1 -2
View File
@@ -1,4 +1,3 @@
# ComfyUI-post-processing-nodes
A collection of post processing nodes for [ComfyUI](https://github.com/comfyanonymous/ComfyUI), simply download this repo and drag the nodes into your
`custom_nodes/` folder
A collection of post processing nodes for [ComfyUI](https://github.com/comfyanonymous/ComfyUI), simply download this repo and drag the nodes into your `custom_nodes/` folder
canny_edge_detect.py
+14 -7
View File
@@ -1,7 +1,8 @@
import numpy as np
import cv2
import numpy as np
import torch
class CannyEdgeDetection:
def __init__(self):
pass
@@ -32,12 +33,18 @@ class CannyEdgeDetection:
CATEGORY = "postprocessing"
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,)
batch_size, height, width, _ = image.shape
result = torch.zeros(batch_size, height, width)
for b in range(batch_size):
tensor_image = image[b].numpy().copy()
gray_image = (cv2.cvtColor(tensor_image, cv2.COLOR_RGB2GRAY) * 255).astype(np.uint8)
canny = cv2.Canny(gray_image, lower_threshold, upper_threshold)
tensor = torch.from_numpy(canny)
result[b] = tensor
return (result,)
NODE_CLASS_MAPPINGS = {
"CannyEdgeDetection": CannyEdgeDetection
}
}
color_correct.py
+42 -36
View File
@@ -1,8 +1,9 @@
import numpy as np
import cv2
import numpy as np
import torch
from PIL import Image, ImageEnhance
class ColorCorrect:
def __init__(self):
pass
@@ -57,52 +58,57 @@ class ColorCorrect:
CATEGORY = "postprocessing"
def color_correct(self, image: torch.Tensor, temperature: float, hue: float, brightness: float, contrast: float, saturation: float, gamma: float):
tensor_image = image.numpy()[0]
batch_size, height, width, _ = image.shape
result = torch.zeros_like(image)
brightness /= 100
contrast /= 100
saturation /= 100
temperature /= 100
for b in range(batch_size):
tensor_image = image[b].numpy()
brightness = 1 + brightness
contrast = 1 + contrast
saturation = 1 + saturation
brightness /= 100
contrast /= 100
saturation /= 100
temperature /= 100
modified_image = Image.fromarray((tensor_image * 255).astype(np.uint8))
brightness = 1 + brightness
contrast = 1 + contrast
saturation = 1 + saturation
# brightness
modified_image = ImageEnhance.Brightness(modified_image).enhance(brightness)
modified_image = Image.fromarray((tensor_image * 255).astype(np.uint8))
# contrast
modified_image = ImageEnhance.Contrast(modified_image).enhance(contrast)
modified_image = np.array(modified_image).astype(np.float32)
# brightness
modified_image = ImageEnhance.Brightness(modified_image).enhance(brightness)
# temperature
if temperature > 0:
modified_image[:, :, 0] *= 1 + temperature
modified_image[:, :, 1] *= 1 + temperature * 0.4
elif temperature < 0:
modified_image[:, :, 2] *= 1 - temperature
modified_image = np.clip(modified_image, 0, 255)/255
# contrast
modified_image = ImageEnhance.Contrast(modified_image).enhance(contrast)
modified_image = np.array(modified_image).astype(np.float32)
# gamma
modified_image = np.clip(np.power(modified_image, gamma), 0, 1)
# temperature
if temperature > 0:
modified_image[:, :, 0] *= 1 + temperature
modified_image[:, :, 1] *= 1 + temperature * 0.4
elif temperature < 0:
modified_image[:, :, 2] *= 1 - temperature
modified_image = np.clip(modified_image, 0, 255)/255
# saturation
hls_img = cv2.cvtColor(modified_image, cv2.COLOR_RGB2HLS)
hls_img[:, :, 2] = np.clip(saturation*hls_img[:, :, 2], 0, 1)
modified_image = cv2.cvtColor(hls_img, cv2.COLOR_HLS2RGB) * 255
# gamma
modified_image = np.clip(np.power(modified_image, gamma), 0, 1)
# 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)
# saturation
hls_img = cv2.cvtColor(modified_image, cv2.COLOR_RGB2HLS)
hls_img[:, :, 2] = np.clip(saturation*hls_img[:, :, 2], 0, 1)
modified_image = cv2.cvtColor(hls_img, cv2.COLOR_HLS2RGB) * 255
modified_image = modified_image.astype(np.uint8)
modified_image = modified_image / 255
modified_image = torch.from_numpy(modified_image).unsqueeze(0)
# 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)
return (modified_image, )
modified_image = modified_image.astype(np.uint8)
modified_image = modified_image / 255
modified_image = torch.from_numpy(modified_image).unsqueeze(0)
result[b] = modified_image
return (result, )
NODE_CLASS_MAPPINGS = {
"ColorCorrect": ColorCorrect,
dither.py
+27 -20
View File
@@ -1,5 +1,6 @@
import torch
class Dither:
def __init__(self):
pass
@@ -24,32 +25,38 @@ class Dither:
CATEGORY = "postprocessing"
def dither(self, image: torch.Tensor, bits: int):
tensor_image = image.numpy()[0]
img = (tensor_image * 255)
height, width, _ = img.shape
batch_size, height, width, _ = image.shape
result = torch.zeros_like(image)
scale = 255 / (2**bits - 1)
for b in range(batch_size):
tensor_image = image[b].numpy()
img = (tensor_image * 255)
height, width, _ = img.shape
for y in range(height):
for x in range(width):
old_pixel = img[y, x].copy()
new_pixel = np.round(old_pixel / scale) * scale
img[y, x] = new_pixel
scale = 255 / (2**bits - 1)
quant_error = old_pixel - new_pixel
for y in range(height):
for x in range(width):
old_pixel = img[y, x].copy()
new_pixel = np.round(old_pixel / scale) * scale
img[y, x] = new_pixel
quant_error = old_pixel - new_pixel
if x + 1 < width:
img[y, x + 1] += quant_error * 7 / 16
if y + 1 < height:
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
img[y, x + 1] += quant_error * 7 / 16
if y + 1 < height:
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
dithered = img / 255
tensor = torch.from_numpy(dithered).unsqueeze(0)
return (tensor,)
dithered = img / 255
tensor = torch.from_numpy(dithered).unsqueeze(0)
result[b] = tensor
return (result,)
NODE_CLASS_MAPPINGS = {
"Dither": Dither
gaussian_blur.py
+12 -5
View File
@@ -1,6 +1,7 @@
import cv2
import torch
class GaussianBlur:
def __init__(self):
pass
@@ -31,11 +32,17 @@ class GaussianBlur:
CATEGORY = "postprocessing"
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,)
batch_size, height, width, _ = image.shape
result = torch.zeros_like(image)
for b in range(batch_size):
tensor_image = image[b].numpy()
blurred = cv2.GaussianBlur(tensor_image, (kernel_size, kernel_size), sigma)
tensor = torch.from_numpy(blurred).unsqueeze(0)
result[b] = tensor
return (result,)
NODE_CLASS_MAPPINGS = {
"GaussianBlur": GaussianBlur
}
}
kmeans_quantize.py
+27 -16
View File
@@ -1,7 +1,8 @@
import numpy as np
import cv2
import numpy as np
import torch
class KMeansQuantize:
def __init__(self):
pass
@@ -32,22 +33,32 @@ class KMeansQuantize:
CATEGORY = "postprocessing"
def kmeans_quantize(self, image: torch.Tensor, colors: int, precision: int):
tensor_image = image.numpy()[0].astype(np.float32)
img = tensor_image
batch_size, height, width, _ = image.shape
result = torch.zeros_like(image)
height, width, c = img.shape
for b in range(batch_size):
tensor_image = image[b].numpy().astype(np.float32)
img = tensor_image
criteria = (
cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER,
precision * 5, 0.01
)
height, width, c = img.shape
img_copy = img.reshape(-1, c)
_, label, center = cv2.kmeans(
img_copy, colors, None,
criteria, 1, cv2.KMEANS_PP_CENTERS
)
criteria = (
cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER,
precision * 5, 0.01
)
result = center[label.flatten()].reshape(*img.shape)
tensor = torch.from_numpy(result).unsqueeze(0)
return (tensor,)
img_copy = img.reshape(-1, c)
_, label, center = cv2.kmeans(
img_copy, colors, None,
criteria, 1, cv2.KMEANS_PP_CENTERS
)
img = center[label.flatten()].reshape(*img.shape)
tensor = torch.from_numpy(img).unsqueeze(0)
result[b] = tensor
return (result,)
NODE_CLASS_MAPPINGS = {
"KMeansQuantize": KMeansQuantize
}
pixel_sort.py
+12 -7
View File
@@ -29,17 +29,22 @@ class PixelSort:
CATEGORY = "postprocessing"
def sort_pixels(self, image, mask, direction, span_limit, sort_by, order):
def sort_pixels(self, image: torch.Tensor, mask: torch.Tensor, direction: str, span_limit: int, sort_by: str, order: str):
horizontal_sort = direction == "horizontal"
reverse_sorting = order == "backward"
sort_by = sort_by[0].upper()
span_limit = span_limit if span_limit > 0 else None
tensor_img = image.numpy()[0]
tensor_mask = mask.numpy()[0]
sorted_image = pixel_sort(tensor_img, tensor_mask, horizontal_sort, span_limit, sort_by, reverse_sorting)
tensor = torch.from_numpy(sorted_image).unsqueeze(0)
return (tensor,)
batch_size = image.shape[0]
result = torch.zeros_like(image)
for b in range(batch_size):
tensor_img = image[b].numpy()
tensor_mask = mask[b].numpy()
sorted_image = pixel_sort(tensor_img, tensor_mask, horizontal_sort, span_limit, sort_by, reverse_sorting)
result[b] = torch.from_numpy(sorted_image)
return (result,)
NODE_CLASS_MAPPINGS = {
"PixelSort": PixelSort,
@@ -132,4 +137,4 @@ def pixel_sort(img, mask, horizontal_sort=False, span_limit=None, sort_by='H', r
sorted_column[start + i] = pixel
sorted_image[:, x] = sorted_column
return sorted_image
return sorted_image
sharpen.py
+18 -11
View File
@@ -1,7 +1,8 @@
import numpy as np
import cv2
import numpy as np
import torch
class Sharpen:
def __init__(self):
pass
@@ -32,19 +33,25 @@ class Sharpen:
CATEGORY = "postprocessing"
def sharpen(self, image: torch.Tensor, kernel_size: int, alpha: float):
tensor_image = image.numpy()[0]
batch_size, height, width, _ = image.shape
result = torch.zeros_like(image)
kernel = np.ones((kernel_size, kernel_size), dtype=np.float32) * -1
center = kernel_size // 2
kernel[center, center] = kernel_size**2
kernel *= alpha
for b in range(batch_size):
tensor_image = image[b].numpy()
sharpened = cv2.filter2D(tensor_image, -1, kernel)
kernel = np.ones((kernel_size, kernel_size), dtype=np.float32) * -1
center = kernel_size // 2
kernel[center, center] = kernel_size**2
kernel *= alpha
tensor = torch.from_numpy(sharpened).unsqueeze(0)
tensor = torch.clamp(tensor, 0, 1)
return (tensor,)
sharpened = cv2.filter2D(tensor_image, -1, kernel)
tensor = torch.from_numpy(sharpened).unsqueeze(0)
tensor = torch.clamp(tensor, 0, 1)
result[b] = tensor
return (result,)
NODE_CLASS_MAPPINGS = {
"Sharpen": Sharpen
}
}