-
Notifications
You must be signed in to change notification settings - Fork 0
/
colorize.py
113 lines (84 loc) · 3.36 KB
/
colorize.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
import os
from PIL import Image
import s3
from werkzeug.utils import secure_filename
import shutil
import subprocess
PIX2PIX_PATH = os.environ.get('PIX2PIX_PATH', 'pix2pix')
def resize_and_align(src_path, dest_path):
"""Resize and double the picture"""
basewidth = 512
baseheight = 512
img = Image.open(src_path)
width, height = img.size
if width >= height:
wpercent = (basewidth/float(width))
hsize = int((float(height)*float(wpercent)))
hsize = hsize - hsize % 4
img = img.resize((basewidth,hsize), Image.ANTIALIAS)
else:
wpercent = (baseheight/float(height))
wsize = int((float(width)*float(wpercent)))
wsize = wsize - wsize % 4
img = img.resize((wsize,baseheight), Image.ANTIALIAS)
img_bw= img.convert('L') #monochromize
aligned_image = Image.new("RGB", (img.size[0] * 2, img.size[1]))
aligned_image.paste(img_bw, (0, 0))
aligned_image.paste(img, (img.size[0], 0))
aligned_image.save(dest_path)
def colorize_photo(filename, model_name):
completed = subprocess.run([
'python', os.path.join(PIX2PIX_PATH, 'test.py'),
'--checkpoints_dir', os.path.join(PIX2PIX_PATH, 'checkpoints'),
'--dataroot', os.path.abspath('uploads'),
'--direction=AtoB',
'--model', 'pix2pix',
'--name', model_name,
'--gpu_ids', '-1',
'--norm', 'instance',
'--image', filename,
], stdout=subprocess.PIPE, check=True)
def resize_back(original_path, processed_path):
img_bw = Image.open(original_path)
img = Image.open(processed_path)
width, height = img_bw.size
img = img.resize((width,height), Image.ANTIALIAS)
img.save(processed_path)
def transfer_color(original_path, colorized_path, processed_path):
subprocess.run([
'convert', colorized_path, original_path,
'-colorspace', 'HSL',
'-compose', 'CopyBlue', '-composite',
'-colorspace', 'sRGB',
processed_path,
], stdout=subprocess.PIPE, check=True)
def process(upload_folder, original_filename, model_name):
"""Returns filename of the processed photo in the uploads folder."""
results_folder = \
os.path.abspath('results/{}/test_latest/images'.format(model_name))
original_path = os.path.abspath(os.path.join(
upload_folder, original_filename))
test_dir = os.path.abspath(os.path.join(
upload_folder, 'test'))
if not os.path.exists(test_dir):
os.makedirs(test_dir)
aligned_path = os.path.join(test_dir, original_filename)
resize_and_align(original_path, aligned_path)
colorize_photo(original_filename, model_name)
basename = original_filename.rsplit('.', 1)[0]
processed_filename = '{}_fake_B.png'.format(basename)
colorized_path = os.path.join(
results_folder, processed_filename)
processed_path = os.path.join(upload_folder, processed_filename)
resize_back(original_path, colorized_path)
transfer_color(original_path, colorized_path, processed_path)
return processed_filename
"""resize_and_align(
os.path.abspath(original_path),
os.path.abspath(target_path),
processed_photo
)
colorize_photo(processed_photo)
shutil.move("path/to/current/file.foo", "path/to/new/destination/for/file.foo")
resize_back(os.path.abspath(original_path), os.path.abspath(processed_photo)
"""