| 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") | |