File size: 1,594 Bytes
ac42b1d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
from sklearn.cluster import KMeans
from collections import Counter
import numpy as np
import cv2
from transformers import pipeline, BasePipeline
class ColorExtractionPipeline(BasePipeline):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.image_pipeline = pipeline("image-classification")
def get_image(self, pil_image):
nimg = np.array(pil_image)
image = cv2.cvtColor(nimg, cv2.COLOR_RGB2BGR)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
return image
def get_labels(self, rimg):
clf = KMeans(n_clusters=5)
labels = clf.fit_predict(rimg)
return labels, clf
def RGB2HEX(self, color):
return "#{:02x}{:02x}{:02x}".format(int(color[0]), int(color[1]), int(color[2]))
def extract_colors(self, pimg):
img = self.get_image(pimg)
reshaped_img = img.reshape(img.shape[0] * img.shape[1], img.shape[2])
labels, clf = self.get_labels(reshaped_img)
counts = Counter(labels)
center_colors = clf.cluster_centers_
ordered_colors = [center_colors[i] for i in counts.keys()]
hex_colors = [self.RGB2HEX(ordered_colors[i]) for i in counts.keys()]
closest_color_to_white = min(center_colors, key=lambda c: np.linalg.norm(c - [255, 255, 255]))
hex_closest_color_to_white = self.RGB2HEX(closest_color_to_white)
return hex_colors, hex_closest_color_to_white
def __call__(self, pimg):
return self.extract_colors(pimg)
color_extraction = ColorExtractionPipeline(task="image-classification")
|