def generate(model, args):
p_wct = model['p_wct']
p_pro = model['p_pro']
# TODO: Use image directly instead of saving to path
content_image_path = '/tmp/content.png'
style_image_path = '/tmp/style.png'
args['content'].save(content_image_path, 'PNG')
args['style'].save(style_image_path, 'PNG')
output_image_path = '/tmp/output.png'
process_stylization.stylization(
stylization_module=p_wct,
smoothing_module=p_pro,
content_image_path=content_image_path,
style_image_path=style_image_path,
# TODO: Allow passing in segmented images
content_seg_path=[],
style_seg_path=[],
output_image_path=output_image_path,
cuda=torch.cuda.is_available(),
save_intermediate=False,
no_post=False)
return {
# TODO: Pass PIL Image directly instead of loading from file
'image': Image.open(output_image_path)
}
def process_image(p_wct,
p_pro,
content_image_path,
content_seg_path=[],
style_image_path='./images/style1.png',
style_seg_path=[],
output_image_path='./results/example1.png',
save_intermediate=False,
no_post=False,
cuda=1,
minsize=256,
maxsize=960):
"""wrapper function of stylization"""
process_stylization.stylization(stylization_module=p_wct,
smoothing_module=p_pro,
content_image_path=content_image_path,
style_image_path=style_image_path,
content_seg_path=content_seg_path,
style_seg_path=style_seg_path,
output_image_path=output_image_path,
cuda=cuda,
save_intermediate=save_intermediate,
no_post=no_post,
minsize=minsize,
maxsize=maxsize)
def stylize(self,
content_image_path,
style_image_path,
output_image_path,
content_seg_path=None,
style_seg_path=None,
smooth=True,
verbose=False):
process_stylization.stylization(
p_wct=self.p_wct,
content_image_path=content_image_path,
style_image_path=style_image_path,
content_seg_path=content_seg_path,
style_seg_path=style_seg_path,
output_image_path=output_image_path,
cuda=self.use_cuda,
smooth=smooth,
verbose=verbose,
)
def main():
parser = argparse.ArgumentParser(
description='Photorealistic Image Stylization')
parser.add_argument(
'--model',
default='./PhotoWCTModels/photo_wct.pth',
help=
'Path to the PhotoWCT model. These are provided by the PhotoWCT submodule, please use `git submodule update --init --recursive` to pull.'
)
parser.add_argument('--content_image_path',
default='./images/content1.png')
parser.add_argument('--content_seg_path', default=[])
parser.add_argument('--style_image_path', default='./images/style1.png')
parser.add_argument('--style_seg_path', default=[])
parser.add_argument('--output_image_path',
default='./results/example1.png')
parser.add_argument('--cuda', type=int, default=1, help='Enable CUDA.')
args = parser.parse_args()
# Load model
p_wct = PhotoWCT()
try:
p_wct.load_state_dict(torch.load(args.model))
except:
print("Fail to load PhotoWCT models. PhotoWCT submodule not updated?")
exit()
if args.cuda:
p_wct.cuda(0)
process_stylization.stylization(
p_wct=p_wct,
content_image_path=args.content_image_path,
style_image_path=args.style_image_path,
content_seg_path=args.content_seg_path,
style_seg_path=args.style_seg_path,
output_image_path=args.output_image_path,
cuda=args.cuda,
)
def run_cyclic_photoWCT(photo_wct_path,
content_image_path,
style_image_path,
stylized_image_path,
reversed_image_path,
cuda,
smoothing_module=None,
do_smoothing=False):
if cuda:
p_wct.load_state_dict(torch.load(photo_wct_path))
else:
p_wct.load_state_dict(
torch.load(photo_wct_path, map_location=torch.device('cpu')))
print("Forward Stylization")
print("Smoothing =", do_smoothing)
# forward stylization
process_stylization.stylization(stylization_module=p_wct,
smoothing_module=smoothing_module,
content_image_path=content_image_path,
style_image_path=style_image_path,
content_seg_path=[],
style_seg_path=[],
output_image_path=stylized_image_path,
cuda=cuda,
do_smoothing=do_smoothing)
print("Reverse Stylization")
print("Smoothing =", do_smoothing)
# reversed stylzation
process_stylization.stylization(stylization_module=p_wct,
smoothing_module=smoothing_module,
content_image_path=stylized_image_path,
style_image_path=content_image_path,
content_seg_path=[],
style_seg_path=[],
output_image_path=reversed_image_path,
cuda=cuda,
do_smoothing=do_smoothing)
styl_seg_folder = os.path.join(folder, 'style_seg')
outp_img_folder = os.path.join(folder, 'results')
cont_img_list = [
f for f in os.listdir(cont_img_folder)
if os.path.isfile(os.path.join(cont_img_folder, f))
]
cont_img_list.sort()
# Load model
p_wct = PhotoWCT()
p_wct.load_state_dict(torch.load(args.model))
p_wct.cuda(0)
for f in cont_img_list:
print("Process " + f)
content_image_path = os.path.join(cont_img_folder, f)
content_seg_path = os.path.join(cont_seg_folder, f).replace(".png", ".pgm")
style_image_path = os.path.join(styl_img_folder, f)
style_seg_path = os.path.join(styl_seg_folder, f).replace(".png", ".pgm")
output_image_path = os.path.join(outp_img_folder, f)
process_stylization.stylization(
p_wct=p_wct,
content_image_path=content_image_path,
style_image_path=style_image_path,
content_seg_path=content_seg_path,
style_seg_path=style_seg_path,
output_image_path=output_image_path,
)
parser.add_argument('--output_image_path', default='./results/example1.png')
parser.add_argument('--save_intermediate', action='store_true', default=False)
parser.add_argument('--fast', action='store_true', default=False)
parser.add_argument('--no_post', action='store_true', default=False)
parser.add_argument('--cuda', type=int, default=1, help='Enable CUDA.')
args = parser.parse_args()
# Load model
p_wct = PhotoWCT()
p_wct.load_state_dict(torch.load(args.model))
if args.fast:
from photo_gif import GIFSmoothing
p_pro = GIFSmoothing(r=35, eps=0.001)
else:
from photo_smooth import Propagator
p_pro = Propagator()
if args.cuda:
p_wct.cuda(0)
process_stylization.stylization(stylization_module=p_wct,
smoothing_module=p_pro,
content_image_path=args.content_image_path,
style_image_path=args.style_image_path,
content_seg_path=args.content_seg_path,
style_seg_path=args.style_seg_path,
output_image_path=args.output_image_path,
cuda=args.cuda,
save_intermediate=args.save_intermediate,
no_post=args.no_post)
def stylize():
content_image_base64 = request.json.get('content_image_base64', None)
if content_image_base64 is None:
raise Exception('content_image_base64 cannot be None')
else:
content_image = base64_png_image_to_pillow_image(content_image_base64)
style_image_base64 = request.json.get('style_image_base64', None)
if style_image_base64 is None:
raise Exception('style_image_base64 cannot be None')
else:
style_image = base64_png_image_to_pillow_image(style_image_base64)
content_segmentation_base64 = request.json.get(
'content_segmentation_base64', None)
if content_segmentation_base64 is None:
content_segmentation = None
else:
content_segmentation = base64_png_image_to_pillow_image(
content_segmentation_base64)
style_segmentation_base64 = request.json.get('style_segmentation_base64',
None)
if style_segmentation_base64 is None:
style_segmentation = None
else:
style_segmentation = base64_png_image_to_pillow_image(
style_segmentation_base64)
if (content_segmentation is None) != (style_segmentation is None):
raise Exception(
'Content segmentation must either be both set or both None')
original_content_image_size = content_image.size
content_image_size = get_apt_image_size(content_image, MAX_NUM_PIXELS)
content_image = content_image.resize(content_image_size, Image.LANCZOS)
if content_segmentation:
content_segmentation = content_segmentation.resize(
content_image_size, Image.LANCZOS)
style_image_size = get_apt_image_size(style_image, MAX_NUM_PIXELS)
style_image = style_image.resize(style_image_size, Image.LANCZOS)
if style_segmentation:
style_segmentation = style_segmentation.resize(style_image_size,
Image.LANCZOS)
content_image_path = get_temp_png_file_path()
content_image.save(content_image_path)
if content_segmentation:
content_segmentation_path = get_temp_png_file_path()
content_segmentation.save(content_segmentation_path)
else:
content_segmentation_path = None
style_image_path = get_temp_png_file_path()
style_image.save(style_image_path)
if style_segmentation:
style_segmentation_path = get_temp_png_file_path()
style_segmentation.save(style_segmentation_path)
else:
style_segmentation_path = None
output_image_path = get_temp_png_file_path()
process_stylization.stylization(
p_wct=p_wct,
content_image_path=content_image_path,
style_image_path=style_image_path,
content_seg_path=content_segmentation_path,
style_seg_path=style_segmentation_path,
output_image_path=output_image_path,
cuda=1,
)
output_image = Image.open(output_image_path)
output_image = output_image.resize(original_content_image_size,
Image.LANCZOS)
image_buffer = six.BytesIO()
output_image.save(image_buffer, format='PNG')
output_image_base64 = base64.b64encode(image_buffer.getvalue())
return jsonify({'output_image_base64': output_image_base64})
cont_img_folder = os.path.join(folder, 'content_img')
cont_seg_folder = os.path.join(folder, 'content_seg')
styl_img_folder = os.path.join(folder, 'style_img')
styl_seg_folder = os.path.join(folder, 'style_seg')
outp_img_folder = os.path.join(folder, 'results')
cont_img_list = [f for f in os.listdir(cont_img_folder) if os.path.isfile(os.path.join(cont_img_folder, f))]
cont_img_list.sort()
# Load model
p_wct = PhotoWCT()
p_wct.load_state_dict(torch.load(args.model))
for f in cont_img_list:
print("Process " + f)
content_image_path = os.path.join(cont_img_folder, f)
content_seg_path = os.path.join(cont_seg_folder, f).replace(".png", ".pgm")
style_image_path = os.path.join(styl_img_folder, f)
style_seg_path = os.path.join(styl_seg_folder, f).replace(".png", ".pgm")
output_image_path = os.path.join(outp_img_folder, f)
process_stylization.stylization(
p_wct=p_wct,
content_image_path=content_image_path,
style_image_path=style_image_path,
content_seg_path=content_seg_path,
style_seg_path=style_seg_path,
output_image_path=output_image_path,
cuda=args.cuda,
)
help=
'Path to the PhotoWCT model. These are provided by the PhotoWCT submodule, please use `git submodule update --init --recursive` to pull.'
)
parser.add_argument('--content_image_path', default='./images/content1.png')
parser.add_argument('--content_seg_path', default=[])
parser.add_argument('--style_image_path', default='./images/style1.png')
parser.add_argument('--style_seg_path', default=[])
parser.add_argument('--output_image_path', default='./results/example1.png')
parser.add_argument('--cuda', type=bool, default=True, help='Enable CUDA.')
args = parser.parse_args()
# Load model
p_wct = PhotoWCT()
try:
p_wct.load_state_dict(torch.load(args.model))
except:
print("Fail to load PhotoWCT models. PhotoWCT submodule not updated?")
exit()
p_wct.cuda(0)
process_stylization.stylization(
p_wct=p_wct,
content_image_path=args.content_image_path,
style_image_path=args.style_image_path,
content_seg_path=args.content_seg_path,
style_seg_path=args.style_seg_path,
output_image_path=args.output_image_path,
cuda=args.cuda,
)
def send_prediction_on_photo(update, context):
global storage
global storage1
# Нам нужно получить две картинки, чтобы произвести перенос стиля, но каждая картинка приходит в
# отдельном апдейте, поэтому в простейшем случае мы будем сохранять id первой картинки в память,
# чтобы, когда уже придет вторая, мы могли загрузить в память уже сами картинки и обработать их.
# Точно место для улучшения, я бы
bot = context.bot
if update.message.text == '/style':
storage.append('1')
update.message.reply_text(
"Пришли 2 фотографии, первая фотография - то, что хочешь изменить. Вторая - стиль, который хочешь получить (любая картина художника)."
)
elif len(storage) != 0:
chat_id = update.message.chat_id
print("Got image from {}".format(chat_id))
# получаем информацию о картинке
image_info = update.message.photo[-1]
image_file = bot.get_file(image_info)
if chat_id in first_image_file:
# первая картинка, которая к нам пришла станет content image, а вторая style image
print("работает style_transfer")
content_image_stream = BytesIO()
first_image_file[chat_id].download(out=content_image_stream)
del first_image_file[chat_id]
style_image_stream = BytesIO()
image_file.download(out=style_image_stream)
output = model.transfer_style(content_image_stream,
style_image_stream,
num_steps=300)
# теперь отправим назад фото
output_stream = BytesIO() #
output.save(output_stream, format='PNG')
output_stream.seek(0)
bot.send_photo(chat_id, photo=output_stream)
print("Sent Photo to user")
storage = []
else:
first_image_file[chat_id] = image_file #
if update.message.text == '/photo_real':
storage1.append('1')
update.message.reply_text(
"Пришли 2 фотографии, первая фотография - пейзаж. Вторая - стиль, который хочешь перенести на свою фотографию."
)
elif len(storage1) != 0:
chat_id = update.message.chat_id
print("Got image from {}".format(chat_id))
# получаем информацию о картинке
image_info = update.message.photo[-1]
image_file = bot.get_file(image_info)
if chat_id in first_image_file:
# первая картинка, которая к нам пришла станет content image, а вторая style image
print("работает photo_real")
content_image_stream = BytesIO()
first_image_file[chat_id].download(out=content_image_stream)
del first_image_file[chat_id]
style_image_stream = BytesIO()
image_file.download(out=style_image_stream)
process_stylization.stylization(
stylization_module=p_wct,
smoothing_module=p_pro,
content_image_path=content_image_stream,
style_image_path=style_image_stream,
content_seg_path=[],
style_seg_path=[],
output_image_path='img.jpg',
cuda=False,
save_intermediate=False,
no_post=False)
output = Image.open('img.jpg')
# теперь отправим назад фото
output_stream = BytesIO() #
output.save(output_stream, format='PNG')
output_stream.seek(0)
bot.send_photo(chat_id, photo=output_stream)
storage1 = []
print("Sent Photo to user")
else:
first_image_file[chat_id] = image_file #
for style_fname in style_fnames:
k_cont = 0
for content_fname in content_fnames:
scene = content_fname.split('/')[0]
output_path = osp.join(CONTENT_PATH, "style_transfer_all", scene, style)
if not osp.isdir(output_path):
os.makedirs(output_path)
out_fname = content_fname.split('/')[-1][:-4] + "___" + style_fname[:-4] + ".png"
if osp.isfile(osp.join(output_path, out_fname)):
k_cont += 1
continue
process_stylization.stylization(p_wct,
p_pro,
osp.join(CONTENT_PATH, "train", content_fname),
osp.join(STYLE_PATH, style, style_fname),
[],
[],
osp.join(output_path, out_fname),
1,
False,
False)
k_cont += 1
print("style: {:s}, sfname: {:s}, {:d}/{:d}".format(style,
style_fname,
k_cont,
len(content_fnames)))
print("Done!")
os.makedirs(args.output_image_path, exist_ok=True)
print("Output folder created!")
# Load model
p_wct = PhotoWCT()
p_wct.load_state_dict(torch.load(args.model_path))
p_pro = Propagator()
p_wct.cuda(0)
style_images_list = os.listdir(args.style_image_path)
style_images_map = {}
content_image = os.path.basename(args.content_image_path)
style_image = random.choice(style_images_list)
style_images_map[content_image] = style_image
stylization(
stylization_module=p_wct,
smoothing_module=p_pro,
content_image_path=args.content_image_path + "/" + "0_s_0_g2.png",
style_image_path=args.style_image_path + "/" + style_image,
content_seg_path=[],
style_seg_path=[],
output_image_path=args.output_image_path,
cuda=1,
save_intermediate=True,
no_post=False)
print("\n{} is done\n".format(args.content_image_path))
print("+++"*20)
with open(args.output_image_path + '/style_images_map.json', 'w') as fp:
json.dump(style_images_map, fp)
