'''PaintTransformer Demo
- 2021-12-21 first created
- See: https://github.com/wzmsltw/PaintTransformer
'''
import os
# import cv2 # <== error
import imageio
import network
from time import time
from glob import glob
from loguru import logger
import numpy as np
import gradio as gr
import paddle
import render_utils
import render_serial
# ---------- Settings ----------
GPU_ID = '-1'
os.environ['CUDA_VISIBLE_DEVICES'] = GPU_ID
DEVICE = 'cpu' if GPU_ID == '-1' else f'cuda:{GPU_ID}'
examples = sorted(glob(os.path.join('input', '*.jpg')))
WIDTH = 512
HEIGHT = 512
STROKE_NUM = 8
FPS = 10
# ---------- Logger ----------
logger.add('app.log', mode='a')
logger.info('===== APP RESTARTED =====')
# ---------- Model ----------
MODEL_FILE = 'paint_best.pdparams'
if not os.path.exists(MODEL_FILE):
os.system('gdown --id 1G0O81qSvGp0kFCgyaQHmPygbVHFi1--q')
logger.info('model downloaded')
else:
logger.info('model already exists')
paddle.set_device(DEVICE)
net_g = network.Painter(5, STROKE_NUM, 256, 8, 3, 3)
net_g.set_state_dict(paddle.load(MODEL_FILE))
net_g.eval()
for param in net_g.parameters():
param.stop_gradient = True
brush_large_vertical = render_utils.read_img('brush/brush_large_vertical.png', 'L')
brush_large_horizontal = render_utils.read_img('brush/brush_large_horizontal.png', 'L')
meta_brushes = paddle.concat([brush_large_vertical, brush_large_horizontal], axis=0)
def predict(image_file):
original_img = render_utils.read_img(image_file, 'RGB', WIDTH, HEIGHT)
logger.info(f'--- image loaded & resized {WIDTH}x{HEIGHT}')
logger.info('--- doing inference...')
t0 = time()
final_result_list = render_serial.render_serial(original_img, net_g, meta_brushes)
logger.info(f'--- inference took {time() - t0:.4f} sec')
# out = cv2.VideoWriter('output.mp4', cv2.VideoWriter_fourcc(*'mp4v'), FPS, (WIDTH, HEIGHT))
frames = []
for idx, frame in enumerate(final_result_list):
# out.write(frame)
frames.append(frame[:,:,::-1]) # BGR -> RGB
# out.release()
imageio.mimsave('output.gif', frames)
logger.info('--- animation generated')
return 'output.gif'
iface = gr.Interface(
predict,
title='🎨 Paint Transformer',
description='This demo converts an image into a sequence of painted images (takes about 2 min ^^;)',
inputs=[
gr.inputs.Image(label='Input image', type='filepath')
],
outputs=[
gr.outputs.Image(label='Output animation')
],
examples=examples,
enable_queue=True,
article='Original work: PaintTransformer
'
)
iface.launch(debug=True)