1 Star 0 Fork 2

X_C/photo2cartoon

forked from linkchainiii/photo2cartoon 
加入 Gitee
与超过 1200万 开发者一起发现、参与优秀开源项目,私有仓库也完全免费 :)
免费加入
文件
克隆/下载
test_onnx.py 2.05 KB
一键复制 编辑 原始数据 按行查看 历史
haoqiang 提交于 2020-12-01 14:38 . upload onnx model
import os
import cv2
import numpy as np
import onnxruntime
import argparse
from utils import Preprocess
parser = argparse.ArgumentParser()
parser.add_argument('--photo_path', type=str, help='input photo path')
parser.add_argument('--save_path', type=str, help='cartoon save path')
args = parser.parse_args()
os.makedirs(os.path.dirname(args.save_path), exist_ok=True)
class Photo2Cartoon:
def __init__(self):
self.pre = Preprocess()
assert os.path.exists('./models/photo2cartoon_weights.onnx'), "[Step1: load weights] Can not find 'photo2cartoon_weights.onnx' in folder 'models!!!'"
self.session = onnxruntime.InferenceSession('./models/photo2cartoon_weights.onnx')
print('[Step1: load weights] success!')
def inference(self, img):
# face alignment and segmentation
face_rgba = self.pre.process(img)
if face_rgba is None:
print('[Step2: face detect] can not detect face!!!')
return None
print('[Step2: face detect] success!')
face_rgba = cv2.resize(face_rgba, (256, 256), interpolation=cv2.INTER_AREA)
face = face_rgba[:, :, :3].copy()
mask = face_rgba[:, :, 3][:, :, np.newaxis].copy() / 255.
face = (face*mask + (1-mask)*255) / 127.5 - 1
face = np.transpose(face[np.newaxis, :, :, :], (0, 3, 1, 2)).astype(np.float32)
# inference
cartoon = self.session.run(['output'], input_feed={'input':face})
# post-process
cartoon = np.transpose(cartoon[0][0], (1, 2, 0))
cartoon = (cartoon + 1) * 127.5
cartoon = (cartoon * mask + 255 * (1 - mask)).astype(np.uint8)
cartoon = cv2.cvtColor(cartoon, cv2.COLOR_RGB2BGR)
print('[Step3: photo to cartoon] success!')
return cartoon
if __name__ == '__main__':
img = cv2.cvtColor(cv2.imread(args.photo_path), cv2.COLOR_BGR2RGB)
c2p = Photo2Cartoon()
cartoon = c2p.inference(img)
if cartoon is not None:
cv2.imwrite(args.save_path, cartoon)
print('Cartoon portrait has been saved successfully!')
马建仓 AI 助手
尝试更多
代码解读
代码找茬
代码优化
Python
1
https://gitee.com/x_c9311/photo2cartoon.git
git@gitee.com:x_c9311/photo2cartoon.git
x_c9311
photo2cartoon
photo2cartoon
master

搜索帮助

0d507c66 1850385 C8b1a773 1850385