|
import os |
|
import cv2 |
|
import numpy as np |
|
import gradio as gr |
|
from inference import Predictor |
|
from utils.image_processing import resize_image |
|
|
|
os.makedirs('output', exist_ok=True) |
|
|
|
|
|
def inference( |
|
image: np.ndarray, |
|
style, |
|
imgsz=None, |
|
retain_color=False, |
|
): |
|
if imgsz is not None: |
|
imgsz = int(imgsz) |
|
|
|
retain_color = retain_color |
|
|
|
weight = { |
|
"AnimeGANv2_Hayao": "GeneratorV2_gldv2_Hayao.pt", |
|
"AnimeGANv2_Shinkai": "GeneratorV2_gldv2_Shinkai.pt", |
|
"AnimeGANv2_Arcane": "GeneratorV2_ffhq_Arcane_210624_e350.pt", |
|
"AnimeGANv2_Test": "GeneratorV2_train_photo_Hayao.pt", |
|
"SummerWar": "GeneratorV2_train_photo_SummerWar.pt", |
|
"Hetalia": "GeneratorV2_train_photo_Hetalia.pt", |
|
}[style] |
|
predictor = Predictor( |
|
weight, |
|
device='cpu', |
|
retain_color=retain_color, |
|
imgsz=imgsz, |
|
) |
|
|
|
save_path = f"output/out.jpg" |
|
image = resize_image(image, width=imgsz) |
|
anime_image = predictor.transform(image)[0] |
|
cv2.imwrite(save_path, anime_image[..., ::-1]) |
|
return anime_image, save_path |
|
|
|
|
|
title = "图片动漫风格转换" |
|
description = r"""将图片转换成动漫风格""" |
|
|
|
gr.Interface( |
|
fn=inference, |
|
inputs=[ |
|
gr.components.Image(label="输入图片"), |
|
gr.Dropdown( |
|
[ |
|
'AnimeGANv2_Hayao', |
|
'AnimeGANv2_Shinkai', |
|
'AnimeGANv2_Arcane', |
|
'AnimeGANv2_Test', |
|
'SummerWar', |
|
'Hetalia', |
|
], |
|
type="value", |
|
value='AnimeGANv2_Hayao', |
|
label='转换风格' |
|
), |
|
gr.Dropdown( |
|
[ |
|
None, |
|
416, |
|
512, |
|
768, |
|
1024, |
|
1536, |
|
], |
|
type="value", |
|
value=None, |
|
label='图片大小' |
|
), |
|
gr.Checkbox(value=False, label="保留原图颜色"), |
|
], |
|
outputs=[ |
|
gr.components.Image(type="numpy", label="转换后图片"), |
|
gr.components.File(label="下载转换图片") |
|
], |
|
title=title, |
|
description=description, |
|
allow_flagging="never", |
|
examples=[ |
|
['example/face/girl4.jpg', 'AnimeGANv2_Arcane', None], |
|
['example/face/leo.jpg', 'AnimeGANv2_Arcane', None], |
|
['example/face/cap.jpg', 'AnimeGANv2_Arcane', None], |
|
['example/face/anne.jpg', 'AnimeGANv2_Arcane', None], |
|
['example/landscape/pexels-camilacarneiro-6318793.jpg', 'AnimeGANv2_Hayao', None], |
|
['example/landscape/pexels-nandhukumar-450441.jpg', 'AnimeGANv2_Hayao', None], |
|
] |
|
).launch() |
|
|
