marker_separation/py/posterize.py

67 lines
1.8 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 = []
original_colors = []
image = None
pallete = None
h = 0
w = 0
n_colors = 3
white = [255, 255, 255]
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 + 1
self.posterize()
self.determine_colors()
def posterize (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)
self.image = bgrquant
cv2.imshow("image", bgrquant)
cv2.waitKey(0)
cv2.destroyAllWindows()
def determine_colors (self):
reshaped = self.image.reshape(-1, self.image.shape[2])
self.original_colors = np.unique(reshaped, axis=0)
2023-10-23 04:01:31 +00:00
#print(self.original_colors)
for i in range(self.n_colors) :
mask = self.extract_color(self.image, self.original_colors[i])
cv2.imwrite(f'{i}.png', mask)
def extract_color (self, image, color):
mask = cv2.inRange(image, color, color)
return cv2.bitwise_not(mask)
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