marker_separation/py/posterize.py

49 lines
1.3 KiB
Python
Raw Normal View History

2023-10-21 12:26:58 +00:00
from sklearn.cluster import MiniBatchKMeans
import numpy as np
import argparse
import cv2
class Posterize:
"""Posterize an image and then find nearest colors to use"""
colors = []
image = None
pallete = None
h = 0
w = 0
n_colors = 3
def __init__ (self, image, pallete, n_colors) :
self.image = cv2.imread(image)
(self.h, self.w) = self.image.shape[:2]
self.pallete = pallete
self.n_colors = n_colors
self.process()
def process (self) :
lab = cv2.cvtColor(self.image, cv2.COLOR_BGR2LAB)
feature = lab.reshape((self.h * self.w, 3))
clusters = MiniBatchKMeans(n_clusters = self.n_colors, n_init = 'auto')
labels = clusters.fit_predict(feature)
quant = clusters.cluster_centers_.astype('uint8')[labels]
2023-10-21 12:26:58 +00:00
rquant = quant.reshape((self.h, self.w, 3))
rfeature = feature.reshape((self.h, self.w, 3))
bgrquant = cv2.cvtColor(rquant, cv2.COLOR_LAB2BGR)
bgrfeature = cv2.cvtColor(rfeature, cv2.COLOR_LAB2BGR)
cv2.imshow("image", bgrquant)
cv2.waitKey(0)
cv2.destroyAllWindows()
2023-10-21 12:26:58 +00:00
def closest(self, colors, color):
colors = np.array(colors)
color = np.array(color)
distances = np.sqrt(np.sum((colors - color) ** 2, axis=1))
index_of_smallest = np.where(distances == np.amin(distances))
smallest_distance = colors[index_of_smallest]
return smallest_distance