Create model.py

model.py

@@ -0,0 +1,41 @@
+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")