latent-space-theories / check_images.py
ludusc's picture
lots of changes
2cafca2
raw
history blame
10.6 kB
#!/usr/bin/env python
import numpy as np
import pandas as pd
from sklearn.svm import SVC
from sklearn.decomposition import PCA
from sklearn.linear_model import LogisticRegression, LinearRegression
from sklearn.model_selection import train_test_split
from tqdm import tqdm
import random
from os.path import join
import os
import pickle
import torch
import matplotlib.pyplot as plt
import PIL
from PIL import Image, ImageColor
import sys
sys.path.append('backend')
from color_annotations import extract_color
from networks_stylegan3 import *
sys.path.append('.')
import dnnlib
import legacy
def hex2rgb(hex_value):
h = hex_value.strip("#")
rgb = tuple(int(h[i:i+2], 16) for i in (0, 2, 4))
return rgb
def rgb2hsv(r, g, b):
# Normalize R, G, B values
r, g, b = r / 255.0, g / 255.0, b / 255.0
# h, s, v = hue, saturation, value
max_rgb = max(r, g, b)
min_rgb = min(r, g, b)
difference = max_rgb-min_rgb
# if max_rgb and max_rgb are equal then h = 0
if max_rgb == min_rgb:
h = 0
# if max_rgb==r then h is computed as follows
elif max_rgb == r:
h = (60 * ((g - b) / difference) + 360) % 360
# if max_rgb==g then compute h as follows
elif max_rgb == g:
h = (60 * ((b - r) / difference) + 120) % 360
# if max_rgb=b then compute h
elif max_rgb == b:
h = (60 * ((r - g) / difference) + 240) % 360
# if max_rgb==zero then s=0
if max_rgb == 0:
s = 0
else:
s = (difference / max_rgb) * 100
# compute v
v = max_rgb * 100
# return rounded values of H, S and V
return tuple(map(round, (h, s, v)))
class DisentanglementBase:
def __init__(self, repo_folder, model, annotations, df, space, colors_list, compute_s=False, variable='H1', categorical=True):
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
print('Using device', self.device)
self.repo_folder = repo_folder
self.model = model.to(self.device)
self.annotations = annotations
self.df = df
self.space = space
self.categorical = categorical
self.variable = variable
self.layers = ['input', 'L0_36_512', 'L1_36_512', 'L2_36_512', 'L3_52_512',
'L4_52_512', 'L5_84_512', 'L6_84_512', 'L7_148_512', 'L8_148_512',
'L9_148_362', 'L10_276_256', 'L11_276_181', 'L12_276_128',
'L13_256_128', 'L14_256_3']
self.layers_shapes = [4, 512, 512, 512, 512, 512, 512, 512, 512, 512, 512, 362, 256, 181, 128, 128]
self.decoding_layers = 16
self.colors_list = colors_list
self.to_hsv()
if compute_s:
self.get_s_space()
def to_hsv(self):
"""
The tohsv function takes the top 3 colors of each image and converts them to HSV values.
It then adds these values as new columns in the dataframe.
:param self: Allow the function to access the dataframe
:return: The dataframe with the new columns added
:doc-author: Trelent
"""
print('Adding HSV encoding')
self.df['H1'] = self.df['top1col'].map(lambda x: rgb2hsv(*hex2rgb(x))[0])
self.df['H2'] = self.df['top2col'].map(lambda x: rgb2hsv(*hex2rgb(x))[0])
self.df['H3'] = self.df['top3col'].map(lambda x: rgb2hsv(*hex2rgb(x))[0])
self.df['S1'] = self.df['top1col'].map(lambda x: rgb2hsv(*hex2rgb(x))[1])
self.df['S2'] = self.df['top2col'].map(lambda x: rgb2hsv(*hex2rgb(x))[1])
self.df['S3'] = self.df['top3col'].map(lambda x: rgb2hsv(*hex2rgb(x))[1])
self.df['V1'] = self.df['top1col'].map(lambda x: rgb2hsv(*hex2rgb(x))[2])
self.df['V2'] = self.df['top2col'].map(lambda x: rgb2hsv(*hex2rgb(x))[2])
self.df['V3'] = self.df['top3col'].map(lambda x: rgb2hsv(*hex2rgb(x))[2])
print('Adding RGB encoding')
self.df['R1'] = self.df['top1col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[0])
self.df['R2'] = self.df['top2col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[0])
self.df['R3'] = self.df['top3col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[0])
self.df['G1'] = self.df['top1col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[1])
self.df['G2'] = self.df['top2col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[1])
self.df['G3'] = self.df['top3col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[1])
self.df['B1'] = self.df['top1col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[2])
self.df['B2'] = self.df['top2col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[2])
self.df['B3'] = self.df['top3col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[2])
return self.df
def get_encoded_latent(self):
# ... (existing code for getX)
if self.space.lower() == 'w':
X = np.array(self.annotations['w_vectors']).reshape((len(self.annotations['w_vectors']), 512))
elif self.space.lower() == 'z':
X = np.array(self.annotations['z_vectors']).reshape((len(self.annotations['z_vectors']), 512))
elif self.space.lower() == 's':
concat_v = []
for i in range(len(self.annotations['w_vectors'])):
concat_v.append(np.concatenate(self.annotations['s_vectors'][i], axis=1))
X = np.array(concat_v)
X = X[:, 0, :]
else:
Exception("Sorry, option not available, select among Z, W, S")
print('Shape embedding:', X.shape)
return X
def get_train_val(self, extremes=False):
X = self.get_encoded_latent()
y = np.array(self.df[self.variable].values)
if self.categorical:
y_cat = pd.cut(y,
bins=[x * 360 / len(self.colors_list) if x < len(self.colors_list)
else 360 for x in range(len(self.colors_list) + 1)],
labels=self.colors_list
).fillna(self.colors_list[0])
x_train, x_val, y_train, y_val = train_test_split(X, y_cat, test_size=0.2)
else:
if extremes:
# Calculate the number of elements to consider (10% of array size)
num_elements = int(0.2 * len(y))
# Get indices of the top num_elements maximum values
top_indices = np.argpartition(array, -num_elements)[-num_elements:]
bottom_indices = np.argpartition(array, -num_elements)[:num_elements]
y_ext = y[top_indices + bottom_indices, :]
X_ext = X[top_indices + bottom_indices, :]
x_train, x_val, y_train, y_val = train_test_split(X_ext, y_ext, test_size=0.2)
else:
x_train, x_val, y_train, y_val = train_test_split(X, y, test_size=0.2)
return x_train, x_val, y_train, y_val
def generate_orig_image(self, vec, seed=False):
"""
The generate_original_image function takes in a latent vector and the model,
and returns an image generated from that latent vector.
:param z: Generate the image
:param model: Generate the image
:return: A pil image
:doc-author: Trelent
"""
G = self.model.to(self.device) # type: ignore
# Labels.
label = torch.zeros([1, G.c_dim], device=self.device)
if seed:
seed = vec
vec = self.annotations['z_vectors'][seed]
Z = torch.from_numpy(vec.copy()).to(self.device)
img = G(Z, label, truncation_psi=1, noise_mode='const')
img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
img = PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB')
return img
def main():
repo_folder = '.'
annotations_file = join(repo_folder, 'data/textile_annotated_files/seeds0000-100000_S.pkl')
with open(annotations_file, 'rb') as f:
annotations = pickle.load(f)
df_file = join(repo_folder, 'data/textile_annotated_files/top_three_colours.csv')
df = pd.read_csv(df_file).fillna('#000000')
model_file = join(repo_folder, 'data/textile_model_files/network-snapshot-005000.pkl')
with dnnlib.util.open_url(model_file) as f:
model = legacy.load_network_pkl(f)['G_ema'] # type: ignore
colors_list = ['Red', 'Orange', 'Yellow', 'Yellow Green', 'Chartreuse Green',
'Kelly Green', 'Green Blue Seafoam', 'Cyan Blue',
'Warm Blue', 'Indigo', 'Purple Magenta', 'Magenta Pink']
colors_list = ['Red Orange', 'Yellow', 'Green', 'Light Blue',
'Blue', 'Purple', 'Pink']
disentanglemnet_exp = DisentanglementBase(repo_folder, model, annotations, df, space='w', colors_list=colors_list)
# x_train, x_val, y_train, y_val = disentanglemnet_exp.get_train_val()
# print(colors_list)
# print(np.unique(y_train, return_counts=True))
# for i, color in enumerate(colors_list):
# idxs = np.where(y_train == color)
# x_color = x_train[idxs][:30, :]
# print(x_color.shape)
# print('Generating images of color ' + color)
# for j, vec in enumerate(x_color):
# vec = np.expand_dims(vec, axis=0)
# img = disentanglemnet_exp.generate_orig_image(vec)
# img.save(f'{repo_folder}/colors_test/color_{color}_{j}.png')
df = disentanglemnet_exp.to_hsv()
df['color'] = pd.cut(df['H1'],
bins=[x * 360 / len(colors_list) if x < len(colors_list)
else 360 for x in range(len(colors_list) + 1)],
labels=colors_list
).fillna(colors_list[0])
print(df['color'].value_counts())
df['seed'] = df['fname'].str.split('/').apply(lambda x: x[-1]).str.replace('seed', '').str.replace('.png','').astype(int)
print(df[df['seed'] == 3][['H1', 'S1', 'V1', 'R1', 'B1', 'G1']])
for i, color in enumerate(colors_list):
idxs = df['color'] == color
x_color = df['seed'][idxs][:30]
print('Generating images of color ' + color)
for j, vec in enumerate(x_color):
img = disentanglemnet_exp.generate_orig_image(int(vec), seed=True)
img.save(f'{repo_folder}/colors_test/color_{color}_{j}corrected.png')
if __name__ == "__main__":
main()