-
Notifications
You must be signed in to change notification settings - Fork 2
/
DuplicateMe.py
190 lines (155 loc) · 6.64 KB
/
DuplicateMe.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
# -*- coding: utf-8 -*-
"""
Created on Wed Aug 5 11:56:28 2020
@author: Vincent Morel
"""
import cv2
import math
import copy
import click
import random
import pickle
import numpy as np
from PIL import Image
from os import path, mkdir
from os.path import isfile, isdir, join, split, splitext, abspath
from datetime import datetime
def Show_Img(img_array):
image = cv2.cvtColor(img_array, cv2.COLOR_BGR2RGB)
Image.fromarray(image).show()
def Get_Colors(src_img,K):
"""
Sample of colors from image using cv2 K-means color extraction.
"""
Z = src_img.reshape((-1,3))
Z = np.float32(Z)
criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)
ret,label,center = cv2.kmeans(Z,K,None,criteria,10,cv2.KMEANS_RANDOM_CENTERS)
# Dominant colors list from kmeans calc
colors = np.uint8(center)
return colors
def Save_to_Disk(logs, img_dir, img):
"""
Save the logs of the image construction and the image itself inside the repository.
"""
base_path, file = split(img_dir)
fname, ext = splitext(file)
folders = {}
folders['logs'] = 'txt'
folders['results'] = 'png'
for folder, ext in folders.items():
folder_dir = str(join(base_path,folder))
output_dir = str(join(base_path,folder,f"{fname}.{ext}"))
if not isdir(folder_dir):
mkdir(folder_dir)
if isfile(output_dir):
now = datetime.now()
dt_string = now.strftime("%H%M%S")
new_fname = fname + dt_string + '.' + ext
output_dir = join(base_path,folder,new_fname)
if folder == 'logs':
with open(output_dir, "wb+") as f:
pickle.dump(logs, f)
print(f"Saved logs at {output_dir}")
else:
cv2.imwrite(output_dir,img)
print(f"Saved image at {output_dir}")
class Canvas:
"""
The Canvas object initializes a blank image as a starting point for the main loop.
Parameters
----------
src_img : (np array)
The input image in a numpy uint8 array.
"""
def __init__(self,src_img,K):
self.h, self.w, self.c = src_img.shape
self.blank_arr = np.full((self.h,self.w,3),255,np.uint8)
self.blank_MSE = np.sqrt((np.subtract(src_img,self.blank_arr,dtype=np.int32))**2).sum()
self.colors = Get_Colors(src_img,K).tolist()
class Parent():
"""
The Parent object represents the parent image to stem from at each generation. The Get_MSE()
method is used calculate the distance between the input image and another candidate. The Mutate()
method is used to apply a random transformation to the parent image.
"""
def __init__(self,img, src_img):
self.img_arr = img
self.src_img = src_img
self.img_arr_copy = copy.deepcopy(self.img_arr)
self.h, self.w, self.c = self.img_arr.shape
self.MSE = self.Get_MSE(self.src_img,self.img_arr)
def Get_MSE(self, src_img, img):
MSE = (np.sqrt((np.subtract(src_img,img,dtype=np.int32))**2).sum())
return MSE
def Mutate(self,N,colors):
# Mutate from parent image
self.img_arr = copy.deepcopy(self.img_arr_copy)
#Initialize random ellipse cv2 variables
clr = random.choice(colors)
rot = random.randint(0,360)
lt = random.randint(0,self.h)
lg = random.randint(0,self.w)
# Size (precision) by generation
s1 = random.randint(0,(self.h//(math.log10((N**6)+20))))
s2 = random.randint(0,(self.w//(math.log10((N**6)+20))))
# draw on img_arr, save metrics
cv2.ellipse(self.img_arr,(lg,lt),(s1,s2),rot,0,360,clr,-1)
self.m_MSE = self.Get_MSE(self.src_img,self.img_arr)
self.m_vars = [clr, rot, lt, lg, s1, s2]
return self.img_arr, self.m_MSE, self.m_vars
@click.command()
@click.argument('img_dir')
@click.argument('k', type=int)
@click.option('--n_generations', type=int)
@click.option('--m_candidates', type=int)
@click.option('--verbose', type=int)
def main(img_dir,k,n_generations=200,m_candidates=100, verbose=0):
"""
This function will start the genetic algorithm process. The program will stem m candidates
from the parent image at each generation. Candidates are created by mutating the parent image.
The mutation is the addition of an ellipse of random coordinates, sizes and color. Once the M
candidates are created for the generation, the program will keep the best mutation and assign it
as the new parent. The best candidate is the one which resembles the most the input image. This
is measured by the MSE. The process loops for N generations.
Parameters
----------
img_dir : (string)
The path of the image to be duplicated using this genetic algorithm.
K : (int)
The number of different colors to sample from during the mutations. Strongly suggested to
adjust this variable in accordance to your image.
n_iterations : (int), optional
The number of generations to evolve from. The default is 200.
n_mutations : (int), optional
The number of mutations to create per generation. The default is 100.
verbose : (bool), optional
A boolean to decide if you want updates on the image being built at every 100th generation.
The default is 0.
"""
src_img = cv2.imread(img_dir)
cvs = Canvas(src_img,k)
darwin_sample = (cvs.blank_arr, cvs.blank_MSE)
darwin_logs = []
try:
for generation in range(n_generations):
p = Parent(darwin_sample[0], src_img)
logs = []
for mutation in range(m_candidates):
m_arr, m_MSE, m_vars = p.Mutate(generation,cvs.colors)
logs.append([m_arr, m_MSE, m_vars])
logs = np.asarray(logs)
MSEs = logs[:,1]
best_mutation_idx = np.where(MSEs == np.amin(MSEs))[0][0]
if np.amin(MSEs) <= darwin_sample[1]:
darwin_sample = (copy.deepcopy(logs[best_mutation_idx][0]), logs[best_mutation_idx][1])
darwin_logs.append([logs[best_mutation_idx,2], logs[best_mutation_idx,1]])
print('MSE :', int(darwin_sample[1]),'\t', f"Progress : {generation+1}/{n_generations}")
if verbose and generation % 100 == 0:
Show_Img(darwin_sample[0])
Save_to_Disk(darwin_logs, img_dir,darwin_sample[0])
Show_Img(darwin_sample[0])
except KeyboardInterrupt:
Save_to_Disk(darwin_logs, img_dir,darwin_sample[0])
if __name__ == "__main__":
main()