├── README.md ├── main.cpp ├── main.py └── testimgs ├── 1.jpg ├── 2.png ├── 3.png └── 4.jpg /README.md: -------------------------------------------------------------------------------- 1 | # photo2cartoon-onnxrun-cpp-py 2 | 使用ONNXRuntime部署StyleGAN人像卡通画,包含C++和Python两个版本的程序 3 | 。起初我是想用opencv做部署的,可是opencv的dnn模块读取onnx文件出错了,无赖只能使用onnxruntime部署。 4 | 本套程序的输入图片是人像大头照,如果输入图片里包含太多背景,需要先做人脸检测+人脸矫正。 5 | 裁剪出人像区域后,作为本套程序的输入图像,否则效果会大打折扣。 6 | 7 | onnx文件,在百度云盘,下载链接:https://pan.baidu.com/s/1FiPHQHZ6VbIyTTBe8qeJXg 8 | 提取码:fdzj 9 | -------------------------------------------------------------------------------- /main.cpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/photo2cartoon-onnxrun-cpp-py/bf53b59ee4cf3f7851c31cfc0020fa5a88138d9c/main.cpp -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | import onnxruntime 4 | import argparse 5 | 6 | class Photo2Cartoon: 7 | def __init__(self): 8 | self.cartoon_net = onnxruntime.InferenceSession('photo2cartoon_models/minivision_female_photo2cartoon.onnx') 9 | self.head_seg_net = onnxruntime.InferenceSession('photo2cartoon_models/minivision_head_seg.onnx') 10 | def inference(self, srcimg): 11 | img = cv2.cvtColor(srcimg, cv2.COLOR_BGR2RGB) 12 | image = cv2.resize(img, (384, 384)) 13 | blob = np.expand_dims(image, axis=0).astype(np.float32) / 255.0 14 | output = self.head_seg_net.run(None, {self.head_seg_net.get_inputs()[0].name: blob})[0].squeeze(axis=0) 15 | 16 | output = cv2.resize(output, (srcimg.shape[1], srcimg.shape[0])) 17 | mask = (output * 255).astype(np.uint8) 18 | 19 | # face_rgba = np.dstack((img, mask)) 20 | # face_rgba = cv2.resize(face_rgba, (256, 256)) 21 | # face = face_rgba[:, :, :3].copy() 22 | # mask = face_rgba[:, :, 3][:, :, np.newaxis].copy() / 255. 23 | 24 | face = cv2.resize(img, (256, 256)) 25 | mask = cv2.resize(mask, (256, 256))[:, :, np.newaxis].copy() / 255. 26 | 27 | face = (face * mask + (1 - mask) * 255) / 127.5 - 1 28 | face = np.transpose(face[np.newaxis, :, :, :], (0, 3, 1, 2)).astype(np.float32) 29 | 30 | # inference 31 | cartoon = self.cartoon_net.run(['output'], input_feed={'input': face}) 32 | 33 | # post-process 34 | cartoon = np.transpose(cartoon[0][0], (1, 2, 0)) 35 | cartoon = (cartoon + 1) * 127.5 36 | cartoon = (cartoon * mask + 255 * (1 - mask)).astype(np.uint8) 37 | cartoon = cv2.cvtColor(cartoon, cv2.COLOR_RGB2BGR) 38 | return cartoon 39 | 40 | if __name__ == '__main__': 41 | parser = argparse.ArgumentParser() 42 | parser.add_argument('--imgpath', type=str, default='testimgs/1.jpg', help='input photo path') 43 | args = parser.parse_args() 44 | 45 | srcimg = cv2.imread(args.imgpath) 46 | c2p = Photo2Cartoon() 47 | cartoon = c2p.inference(srcimg) 48 | 49 | cv2.namedWindow('srcimg', cv2.WINDOW_NORMAL) 50 | cv2.imshow('srcimg', srcimg) 51 | winName = 'Deep learning object detection in ONNXRuntime' 52 | cv2.namedWindow(winName, cv2.WINDOW_NORMAL) 53 | cv2.imshow(winName, cartoon) 54 | cv2.waitKey(0) 55 | cv2.destroyAllWindows() -------------------------------------------------------------------------------- /testimgs/1.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/photo2cartoon-onnxrun-cpp-py/bf53b59ee4cf3f7851c31cfc0020fa5a88138d9c/testimgs/1.jpg -------------------------------------------------------------------------------- /testimgs/2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/photo2cartoon-onnxrun-cpp-py/bf53b59ee4cf3f7851c31cfc0020fa5a88138d9c/testimgs/2.png -------------------------------------------------------------------------------- /testimgs/3.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/photo2cartoon-onnxrun-cpp-py/bf53b59ee4cf3f7851c31cfc0020fa5a88138d9c/testimgs/3.png -------------------------------------------------------------------------------- /testimgs/4.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/photo2cartoon-onnxrun-cpp-py/bf53b59ee4cf3f7851c31cfc0020fa5a88138d9c/testimgs/4.jpg --------------------------------------------------------------------------------