forked from Onr/Council-GAN
-
Notifications
You must be signed in to change notification settings - Fork 1
/
main.py
416 lines (371 loc) · 16 KB
/
main.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
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
# -*- coding: utf-8 -*-
from flask import Flask, render_template, request, redirect, url_for, jsonify, Response, after_this_request, send_file
from flask_cors import CORS
from werkzeug.utils import secure_filename
from test_on_folder import runImageTransfer
from utils import get_config, get_data_loader_folder, pytorch03_to_pytorch04
from trainer_council import Council_Trainer
from torch import nn
from scipy.stats import entropy
import torch.nn.functional as F
from torch.autograd import Variable
from data import ImageFolder
import numpy as np
import torchvision.utils as vutils
try:
from itertools import izip as zip
except ImportError: # will be 3.x series
pass
import torch
import os
import random
import string
import uuid
import shutil
from queue import Queue, Empty
from threading import Thread
from tqdm import tqdm
import PIL
from PIL import Image, ImageOps
import base64
import io
import signal
import sys
app = Flask(__name__)
app.config['JSON_AS_ASCII'] = False
CORS(app)
path = "./static"
threads = []
#multi-threads with return value
class ThreadWithReturnValue(Thread):
def __init__(self, group=None, target=None, name=None,
args=(), kwargs={}, Verbose=None):
Thread.__init__(self, group, target, name, args, kwargs)
self._return = None
def run(self):
print(type(self._target))
if self._target is not None:
self._return = self._target(*self._args,
**self._kwargs)
def join(self, timeout=3):
Thread.join(self, timeout=3)
return self._return
class thread_with_trace(ThreadWithReturnValue):
def __init__(self, *args, **keywords):
ThreadWithReturnValue.__init__(self, *args, **keywords)
self.killed = False
def start(self):
self.__run_backup = self.run
self.run = self.__run
ThreadWithReturnValue.start(self)
def __run(self):
sys.settrace(self.globaltrace)
self.__run_backup()
self.run = self.__run_backup
def globaltrace(self, frame, event, arg):
if event == 'call':
return self.localtrace
else:
return None
def localtrace(self, frame, event, arg):
if self.killed:
if event == 'line':
raise SystemExit()
return self.localtrace
def kill(self):
self.killed = True
############################################################
#preload model
def loadModel(config, checkpoint, a2b):
seed = 1
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
# Load experiment setting
config = get_config(config)
input_dim = config['input_dim_a'] if a2b else config['input_dim_b']
council_size = config['council']['council_size']
style_dim = config['gen']['style_dim']
trainer = Council_Trainer(config)
only_one = False
if 'gen_' in checkpoint[-21:]:
state_dict = torch.load(checkpoint)
try:
print(state_dict)
if a2b:
trainer.gen_a2b_s[0].load_state_dict(state_dict['a2b'])
else:
trainer.gen_b2a_s[0].load_state_dict(state_dict['b2a'])
except:
print('a2b should be set to ' + str(not a2b) + ' , Or config file could be wrong')
a2b = not a2b
if a2b:
trainer.gen_a2b_s[0].load_state_dict(state_dict['a2b'])
else:
trainer.gen_b2a_s[0].load_state_dict(state_dict['b2a'])
council_size = 1
only_one = True
else:
for i in range(council_size):
try:
if a2b:
tmp_checkpoint = checkpoint[:-8] + 'a2b_gen_' + str(i) + '_' + checkpoint[-8:] + '.pt'
state_dict = torch.load(tmp_checkpoint)
trainer.gen_a2b_s[i].load_state_dict(state_dict['a2b'])
else:
tmp_checkpoint = checkpoint[:-8] + 'b2a_gen_' + str(i) + '_' + checkpoint[-8:] + '.pt'
state_dict = torch.load(tmp_checkpoint)
trainer.gen_b2a_s[i].load_state_dict(state_dict['b2a'])
except:
print('a2b should be set to ' + str(not a2b) + ' , Or config file could be wrong')
a2b = not a2b
if a2b:
tmp_checkpoint = checkpoint[:-8] + 'a2b_gen_' + str(i) + '_' + checkpoint[-8:] + '.pt'
state_dict = torch.load(tmp_checkpoint)
trainer.gen_a2b_s[i].load_state_dict(state_dict['a2b'])
else:
tmp_checkpoint = checkpoint[:-8] + 'b2a_gen_' + str(i) + '_' + checkpoint[-8:] + '.pt'
state_dict = torch.load(tmp_checkpoint)
trainer.gen_b2a_s[i].load_state_dict(state_dict['b2a'])
trainer.cuda()
trainer.eval()
return [trainer, config, council_size, style_dim]
peson2anime_preloadModel = loadModel("pretrain/anime/256/anime2face_council_folder.yaml", "pretrain/anime/256/01000000", 0)
male2female_preloadModel = loadModel("pretrain/m2f/256/male2female_council_folder.yaml", "pretrain/m2f/256/01000000", 1)
noglasses_preloadModel = loadModel("pretrain/glasses_removal/128/glasses_council_folder.yaml", "pretrain/glasses_removal/128/01000000", 1)
############################################################
# 업로드 HTML 렌더링
@app.route('/')
def render_file():
return render_template('upload.html')
@app.route('/healthz', methods=['GET'])
def healthz():
return "I am alive", 200
#for web Demo
@app.route('/fileUpload', methods=['POST'])
def fileupload():
#내가 전달 받는 request는 'file'과 'check_model'
check_value = request.form['check_model']
f = request.files['file']
is_api = False
#handling over request by status code 429
global threads
print(len(threads),"fileUpload")
if len(threads) > 10:
return Response("error : Too many requests", status=429)
try:
randomDirName = str(uuid.uuid4()) #사용자끼리의 업로드한 이미지가 겹치지 않게끔 uuid를 이용하여 사용자를 구분하는 디렉터리를 만든다.
if check_value == "ani":
path = '/home/user/upload/person2anime/'
os.mkdir(path + randomDirName)
f.save(path + randomDirName + '/' + secure_filename(f.filename))
byte_image_list = person_To_anime(randomDirName, is_api)
return render_template('showImage.html', rawimg=byte_image_list)
elif check_value == "m2f":
path = '/home/user/upload/male2female/'
os.mkdir(path + randomDirName)
f.save(path + randomDirName + '/' +secure_filename(f.filename))
byte_image_list = male_To_female(randomDirName, is_api)
return render_template('showImage.html', rawimg=byte_image_list)
else:
path = '/home/user/upload/no_glasses/'
os.mkdir(path + randomDirName)
f.save(path + randomDirName + '/' + secure_filename(f.filename))
byte_image_list = no_glasses(randomDirName, is_api)
return render_template('showImage.html', rawimg=byte_image_list)
except Exception as e:
print(e)
return Response("upload file and load model is fail", status=400)
#for swagger file
@app.route('/convert_image', methods=['POST'])
def convert_imgae():
check_value = request.form['check_model']
f = request.files['file']
is_api = True
#handling over request by status code 429
global threads
print(len(threads),"convert_image")
if len(threads) > 10:
return Response("error : Too many requests", status=429)
try:
randomDirName = str(uuid.uuid4()) #사용자끼리의 업로드한 이미지가 겹치지 않게끔 uuid를 이용하여 사용자를 구분하는 디렉터리를 만든다.
if check_value == "ani":
path = '/home/user/upload/person2anime/'
os.mkdir(path + randomDirName)
f.save(path + randomDirName + '/' + secure_filename(f.filename))
byte_image_list = person_To_anime(randomDirName, is_api)
return send_file(byte_image_list, mimetype="image/jpeg")
elif check_value == "m2f":
path = '/home/user/upload/male2female/'
os.mkdir(path + randomDirName)
f.save(path + randomDirName + '/' +secure_filename(f.filename))
byte_image_list = male_To_female(randomDirName, is_api)
return send_file(byte_image_list, mimetype="image/jpeg")
else:
path = '/home/user/upload/no_glasses/'
os.mkdir(path + randomDirName)
f.save(path + randomDirName + '/' + secure_filename(f.filename))
byte_image_list = no_glasses(randomDirName, is_api)
return send_file(byte_image_list, mimetype="image/jpeg")
except Exception as e:
print(e)
return Response("upload file and load model is fail", status=400)
#사용자의 입력을 받아서 각 원하는 결과물을 라우팅
def person_To_anime(randomDirName, is_api):
try:
user_key = randomDirName
input_ = "/home/user/upload/person2anime/" + user_key
a2b = 0
model_type = 'person2anime'
#handling multi-threads
t1 = thread_with_trace(target=runImageTransfer, args=(peson2anime_preloadModel,input_,user_key,a2b))
t1.user_id = user_key
threads.append(t1)
while threads[0].user_id!=user_key:
print(str(user_key)+": ", threads[0].user_id)
if threads[0].is_alive():
threads[0].join()
threads[0].start()
file_list = threads[0].join(timeout=3)
if threads[0].is_alive():
threads[0].kill()
threads.pop(0)
raise Exception("error model does not work! please try again 30 seconds later")
threads.pop(0)
file_list.sort()
byte_image_list = [] #byte_image를 담기위한 list
tmp_list = [] #byte_image를 담기전에 decode 하기 위한 list
#imgFIle은 np.array형태여야 fromarray에 담길수 있음
#img_io는 각 파일마다 byte 객체를 적용해줘야하므로 for문 안에서 같이 반복을 돌아야 함
#b64encode로 encode 해준다.
for image in file_list:
imgFile = PIL.Image.fromarray(np.array(PIL.Image.open(image).convert("RGB")))
img_io = io.BytesIO()
imgFile.save(img_io, 'jpeg', quality = 100)
img_io.seek(0)
if is_api : return img_io
img = base64.b64encode(img_io.getvalue())
tmp_list.append(img)
#decode 작업은 여기서 해준다.
for i in tmp_list:
byte_image_list.append(i.decode('ascii'))
#input file과 output file을 모두 제거해주는 함수 호출
remove(user_key, model_type)
return byte_image_list
#return render_template('showImage.html', rawimg=byte_image_list)
except Exception as e:
print(e)
return Response("person2anime is fail", status=400)
def male_To_female(randomDirName, is_api):
try:
user_key = randomDirName
input_ = "/home/user/upload/male2female/" + user_key
a2b = 1
model_type = 'male2female'
t1 = thread_with_trace(target=runImageTransfer, args=(male2female_preloadModel,input_,user_key,a2b))
t1.user_id = user_key
threads.append(t1)
while threads[0].user_id!=user_key:
print(str(user_key)+": ", threads[0].user_id)
if threads[0].is_alive():
threads[0].join()
threads[0].start()
file_list = threads[0].join(timeout=3)
if threads[0].is_alive():
threads[0].kill()
threads.pop(0)
raise Exception("error model does not work! please try again 30 seconds later")
threads.pop(0)
file_list.sort()
byte_image_list = [] #byte_image를 담기위한 list
tmp_list = [] #byte_image를 담기전에 decode 하기 위한 list
#imgFIle은 np.array형태여야 fromarray에 담길수 있음
#img_io는 각 파일마다 byte 객체를 적용해줘야하므로 for문 안에서 같이 반복을 돌아야 함
#b64encode로 encode 해준다.
for image in file_list:
imgFile = PIL.Image.fromarray(np.array(PIL.Image.open(image).convert("RGB")))
img_io = io.BytesIO()
imgFile.save(img_io, 'jpeg', quality = 100)
img_io.seek(0)
if is_api : return img_io
img = base64.b64encode(img_io.getvalue())
tmp_list.append(img)
#decode 작업은 여기서 해준다.
for i in tmp_list:
byte_image_list.append(i.decode('ascii'))
#input file과 output file을 모두 제거해주는 함수 호출
remove(user_key, model_type)
return byte_image_list
#return render_template('showImage.html', rawimg=byte_image_list)
except Exception as e:
print(e)
return Response("male2female is fail", status=400)
def no_glasses(randomDirName, is_api):
try:
user_key = randomDirName
input_ = "/home/user/upload/no_glasses/" + user_key
a2b = 1
model_type = 'no_glasses'
t1 = thread_with_trace(target=runImageTransfer, args=(noglasses_preloadModel,input_,user_key,a2b))
t1.user_id = user_key
threads.append(t1)
while threads[0].user_id!=user_key:
print(str(user_key)+": ", threads[0].user_id)
if threads[0].is_alive():
threads[0].join()
threads[0].start()
file_list = threads[0].join(timeout=3)
if threads[0].is_alive():
threads[0].kill()
threads.pop(0)
raise Exception("error model does not work! please try again 30 seconds later")
threads.pop(0)
file_list.sort()
byte_image_list = [] #byte_image를 담기위한 list
tmp_list = [] #byte_image를 담기전에 decode 하기 위한 list
#imgFIle은 np.array형태여야 fromarray에 담길수 있음
#img_io는 각 파일마다 byte 객체를 적용해줘야하므로 for문 안에서 같이 반복을 돌아야 함
#b64encode로 encode 해준다.
for image in file_list:
imgFile = PIL.Image.fromarray(np.array(PIL.Image.open(image).convert("RGB")))
img_io = io.BytesIO()
imgFile.save(img_io, 'jpeg', quality = 100)
img_io.seek(0)
if is_api : return img_io
img = base64.b64encode(img_io.getvalue())
tmp_list.append(img)
#decode 작업은 여기서 해준다.
for i in tmp_list:
byte_image_list.append(i.decode('ascii'))
#input file과 output file을 모두 제거해주는 함수 호출
remove(user_key, model_type)
return byte_image_list
#return render_template('showImage.html', rawimg=byte_image_list)
except Exception as e:
print(e)
return Response("no_glasses is fail", status=400)
def remove(user_key, model_type):
remove_input_dir = '/home/user/upload/' + model_type + '/' + user_key
path = os.path.join('static/img/', user_key)
print("Now start to remove file")
print("user key is " + user_key)
print("Input path " + remove_input_dir)
print("Output path " + path)
#output path를 삭제하는 try 문
try:
if os.path.isdir(path):
shutil.rmtree(path)
print("Delete " + path + " is completed")
except Exception as e:
print(e)
print("Delete" + path + " is failed")
#input path를 삭제하는 try 문
try:
if os.path.isdir(remove_input_dir):
shutil.rmtree(remove_input_dir)
print("Delete" + remove_input_dir + " is completed")
except Exception as e:
print("Delete" + remove_input_dir + " is failed")
return print("All of delete process is completed!")
if __name__ == '__main__':
# server execute
app.run(host='0.0.0.0', port=80)