Upload 3 files

extract-colors.py

@@ -0,0 +1,37 @@
+from sklearn.cluster import KMeans
+from collections import Counter
+import numpy as np
+import cv2
+
+def get_image(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(rimg):
+    clf = KMeans(n_clusters=5)
+    labels = clf.fit_predict(rimg)
+    return labels, clf
+
+def get_closest_color(colors):
+    white = (255, 255, 255)
+    closest_color = min(colors, key=lambda c: np.linalg.norm(np.array(c) - white))
+    return closest_color
+
+def RGB2HEX(color):
+    return "#{:02x}{:02x}{:02x}".format(int(color[0]), int(color[1]), int(color[2]))
+
+def extract_colors_and_closest_to_white(image_path):
+    img = cv2.imread(image_path)
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    reshaped_img = img.reshape(img.shape[0] * img.shape[1], img.shape[2])
+    labels, clf = get_labels(reshaped_img)
+    counts = Counter(labels)
+    center_colors = clf.cluster_centers_
+    ordered_colors = [center_colors[i] for i in counts.keys()]
+    hex_colors = [RGB2HEX(ordered_colors[i]) for i in counts.keys()]
+
+    closest_color_to_white = get_closest_color(center_colors)
+    hex_closest_color_to_white = RGB2HEX(closest_color_to_white)
+    return hex_colors, hex_closest_color_to_white

predict.py

@@ -0,0 +1,17 @@
+from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
+from PIL import Image
+from io import BytesIO
+import requests
+import numpy as np
+
+model = AutoModelForSequenceClassification.from_pretrained("./model/extract-colors.py")
+tokenizer = AutoTokenizer.from_pretrained("./model/tokenizer", tokenizer_func=tokenizer_function)
+model_class = pipeline('image-classification', model=model, tokenizer=tokenizer)
+
+def get_colors_and_closest_to_white(image_url):
+    response = requests.get(image_url)
+    img = Image.open(BytesIO(response.content))
+    img_array = np.array(img)
+
+    result = model_class(image_array=img_array)
+    return result

tokenizer.py

@@ -0,0 +1,2 @@
+def tokenizer_function(text):
+    return text.strip().split('\n')