65 lines
1.7 KiB
Python
65 lines
1.7 KiB
Python
import cv2
|
|
import numpy as np
|
|
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):
|
|
batch_size, height, width, _ = image.shape
|
|
result = torch.zeros_like(image)
|
|
|
|
for b in range(batch_size):
|
|
tensor_image = image[b].numpy().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
|
|
)
|
|
|
|
img = center[label.flatten()].reshape(*img.shape)
|
|
tensor = torch.from_numpy(img).unsqueeze(0)
|
|
result[b] = tensor
|
|
|
|
return (result,)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"KMeansQuantize": KMeansQuantize
|
|
}
|