marker_separation/py/posterize.py

60 lines
1.6 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
2023-10-24 20:45:22 +00:00
from common import convert_color, closest_color
2023-10-21 12:26:58 +00:00
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_mask(self.image, self.original_colors[i])
cv2.imwrite(f'{i}.png', mask)
def extract_color_mask (self, image, color):
mask = cv2.inRange(image, color, color)
return cv2.bitwise_not(mask)