latent-space-theories / backend /color_annotations.py
ludusc's picture
created global cleaned file
973a4da
#!/usr/bin/env python
"""Extract color features from the generated textile images."""
import os
#os.environ["TOKENIZERS_PARALLELISM"] = "false"
from tqdm import tqdm
#from transformers import pipeline
import numpy as np
import pandas as pd
import time
import click
from PIL import Image
import math
import pickle
from glob import glob
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import matplotlib.image as mpimg
import cv2
import extcolors
from colormap import rgb2hex
from PIL import Image
from matplotlib.offsetbox import OffsetImage, AnnotationBbox
def color_to_df(input):
colors_pre_list = str(input).replace('([(','').split(', (')[0:-1]
df_rgb = [i.split('), ')[0] + ')' for i in colors_pre_list]
df_percent = [i.split('), ')[1].replace(')','') for i in colors_pre_list]
#convert RGB to HEX code
df_color_up = [rgb2hex(int(i.split(", ")[0].replace("(","")),
int(i.split(", ")[1]),
int(i.split(", ")[2].replace(")",""))) for i in df_rgb]
df = pd.DataFrame(zip(df_color_up, df_percent), columns = ['c_code','occurence'])
return df
def extract_color(input_image, tolerance, zoom, outpath, save=None):
colors_x = extcolors.extract_from_image(input_image, tolerance = tolerance, limit = 13)
df_color = color_to_df(colors_x)
#annotate text
list_color = list(df_color['c_code'])
list_precent = [int(i) for i in list(df_color['occurence'])]
text_c = [c + ' ' + str(round(p*100/sum(list_precent),1)) +'%' for c, p in zip(list_color, list_precent)]
colors = list(df_color['c_code'])
if '#000000' in colors:
colors.remove('#000000')
return colors[:3]
@click.command()
@click.option('--genimages_dir', help='Where the output images are saved', type=str, required=True, metavar='DIR')
def annotate_textile_images(
genimages_dir: str,
):
"""Produce annotations for the generated images.
\b
#
python annotate_textiles.py --genimages_dir /home/ludosc/data/stylegan-10000-textile-upscale
"""
colours = []
pickle_files = glob(genimages_dir + '/imgs0000*.pkl')
for pickle_file in pickle_files:
print('Using pickle file: ', pickle_file)
with open(pickle_file, 'rb') as f:
info = pickle.load(f)
listlen = len(info['fname'])
os.makedirs('/data/ludosc/colour_palettes/', exist_ok=True)
for i,im in enumerate(tqdm(info['fname'])):
try:
top_cols = exact_color(im, 12, 5, '/data/ludosc/colour_palettes/' + im.split('/')[-1])
colours.append([im]+top_cols)
except Exception as e:
print(e)
if i % 1000 == 0:
df = pd.DataFrame(colours, columns=['fname', 'top1col', 'top2col', 'top3col'])
print(df.head())
df.to_csv(genimages_dir + f'/top_three_colours.csv', index=False)
df = pd.DataFrame(colours, columns=['fname', 'top1col', 'top2col', 'top3col'])
print(df.head())
df.to_csv(genimages_dir + f'/final_sim_{os.path.basename(pickle_file.split(".")[0])}.csv', index=False)
#----------------------------------------------------------------------------
if __name__ == "__main__":
annotate_textile_images() # pylint: disable=no-value-for-parameter
#----------------------------------------------------------------------------