testaus3 / extract.py
ismot's picture
Duplicate from huggingface-projects/color-palette-generator-sd
4e0a699
from colorthief import ColorThief
from pathlib import Path
import json
from PIL import Image
images_path = Path('frontend/static/images')
images = images_path.glob("*.[jpeg jpg png]*")
print(images)
data = {}
for image in images:
print(image.stem)
image_pil = Image.open(image)
color_thief = ColorThief(image)
image_pil.save(Path.joinpath(images_path, (image.stem + ".jpg")), optimize=True, quality=95)
prompt = image.stem.split("-")[2]
try:
type(data[prompt]) == list
except:
data[prompt] = []
colors = color_thief.get_palette(color_count=5, quality=1)
colors_hex = ['#%02x%02x%02x' % (color) for color in colors]
data[prompt].append({
"colors": colors_hex,
"imgURL": "static/images/" + image.stem + ".jpg"
})
prompts = [{"prompt": prompt, "images": values}
for (prompt, values) in data.items()]
with open('frontend/static/data.json', 'w') as f:
json.dump(prompts, f)