From 54364d4ed4e3537ec5dac38a96ecbdcadc639318 Mon Sep 17 00:00:00 2001 From: jm12138 <2286040843@qq.com> Date: Mon, 16 Nov 2020 11:42:28 +0800 Subject: [PATCH] Add the UGATIT Modules --- .../style_transfer/UGATIT_100w/README.md | 122 ++++++++++++++++++ .../image/style_transfer/UGATIT_100w/model.py | 67 ++++++++++ .../style_transfer/UGATIT_100w/module.py | 69 ++++++++++ .../style_transfer/UGATIT_100w/processor.py | 120 +++++++++++++++++ .../image/style_transfer/UGATIT_83w/README.md | 122 ++++++++++++++++++ .../image/style_transfer/UGATIT_83w/model.py | 67 ++++++++++ .../image/style_transfer/UGATIT_83w/module.py | 69 ++++++++++ .../style_transfer/UGATIT_83w/processor.py | 120 +++++++++++++++++ .../image/style_transfer/UGATIT_92w/README.md | 122 ++++++++++++++++++ .../image/style_transfer/UGATIT_92w/model.py | 67 ++++++++++ .../image/style_transfer/UGATIT_92w/module.py | 69 ++++++++++ .../style_transfer/UGATIT_92w/processor.py | 120 +++++++++++++++++ 12 files changed, 1134 insertions(+) create mode 100644 hub_module/modules/image/style_transfer/UGATIT_100w/README.md create mode 100644 hub_module/modules/image/style_transfer/UGATIT_100w/model.py create mode 100644 hub_module/modules/image/style_transfer/UGATIT_100w/module.py create mode 100644 hub_module/modules/image/style_transfer/UGATIT_100w/processor.py create mode 100644 hub_module/modules/image/style_transfer/UGATIT_83w/README.md create mode 100644 hub_module/modules/image/style_transfer/UGATIT_83w/model.py create mode 100644 hub_module/modules/image/style_transfer/UGATIT_83w/module.py create mode 100644 hub_module/modules/image/style_transfer/UGATIT_83w/processor.py create mode 100644 hub_module/modules/image/style_transfer/UGATIT_92w/README.md create mode 100644 hub_module/modules/image/style_transfer/UGATIT_92w/model.py create mode 100644 hub_module/modules/image/style_transfer/UGATIT_92w/module.py create mode 100644 hub_module/modules/image/style_transfer/UGATIT_92w/processor.py diff --git a/hub_module/modules/image/style_transfer/UGATIT_100w/README.md b/hub_module/modules/image/style_transfer/UGATIT_100w/README.md new file mode 100644 index 00000000..a66abda1 --- /dev/null +++ b/hub_module/modules/image/style_transfer/UGATIT_100w/README.md @@ -0,0 +1,122 @@ +## 模型概述 +UGATIT 图像风格转换模型 + +模型可将输入的人脸图像转换成动漫风格 + +模型权重来自UGATIT-Paddle开源项目 + +模型所使用的权重为genA2B_1000000 + +模型详情请参考[UGATIT-Paddle开源项目](https://github.com/miraiwk/UGATIT-paddle) + +## 模型安装 + +```shell +$hub install UGATIT_100w +``` + + +## API 说明 + +```python +def style_transfer( + self, + images=None, + paths=None, + batch_size=1, + output_dir='output', + visualization=False +) +``` + +风格转换API，将输入的人脸图像转换成动漫风格。 + +转换效果图如下： + +![输入图像](https://ai-studio-static-online.cdn.bcebos.com/d130fabd8bd34e53b2f942b3766eb6bbd3c19c0676d04abfbd5cc4b83b66f8b6) +![输出图像](https://ai-studio-static-online.cdn.bcebos.com/8538af03b3f14b1884fcf4eec48965baf939e35a783d40129085102057438c77) + +**参数** + +* images (list\[numpy.ndarray\]): 图片数据，ndarray.shape 为 \[H, W, C\]，默认为 None； +* paths (list\[str\]): 图片的路径，默认为 None； +* batch\_size (int): batch 的大小，默认设为 1； +* visualization (bool): 是否将识别结果保存为图片文件，默认设为 False； +* output\_dir (str): 图片的保存路径，默认设为 output。 + + +**返回** + +* res (list\[numpy.ndarray\]): 输出图像数据，ndarray.shape 为 \[H, W, C\]。 + + +## 预测代码示例 + +```python +import cv2 +import paddlehub as hub + +# 模型加载 +# use_gpu：是否使用GPU进行预测 +model = hub.Module('UGATIT_100w', use_gpu=False) + +# 模型预测 +result = model.style_transfer(images=[cv2.imread('/PATH/TO/IMAGE')]) + +# or +# result = model.style_transfer(paths=['/PATH/TO/IMAGE']) +``` + +## 服务部署 + +PaddleHub Serving可以部署一个在线图像风格转换服务。 + +## 第一步：启动PaddleHub Serving + +运行启动命令： +```shell +$ hub serving start -m UGATIT_100w +``` + +这样就完成了一个图像风格转换的在线服务API的部署，默认端口号为8866。 + +**NOTE:** 如使用GPU预测，则需要在启动服务之前，请设置CUDA\_VISIBLE\_DEVICES环境变量，否则不用设置。 + +## 第二步：发送预测请求 + +配置好服务端，以下数行代码即可实现发送预测请求，获取预测结果 + +```python +import requests +import json +import cv2 +import base64 + + +def cv2_to_base64(image): + data = cv2.imencode('.jpg', image)[1] + return base64.b64encode(data.tostring()).decode('utf8') + + +# 发送HTTP请求 +data = {'images':[cv2_to_base64(cv2.imread("/PATH/TO/IMAGE"))]} +headers = {"Content-type": "application/json"} +url = "http://127.0.0.1:8866/predict/UGATIT_100w" +r = requests.post(url=url, headers=headers, data=json.dumps(data)) + +# 打印预测结果 +print(r.json()["results"]) +``` + + +## 模型相关信息 + +### 模型代码 + +https://github.com/miraiwk/UGATIT-paddle + +### 依赖 + +paddlepaddle >= 1.8.0 + +paddlehub >= 1.8.0 \ No newline at end of file diff --git a/hub_module/modules/image/style_transfer/UGATIT_100w/model.py b/hub_module/modules/image/style_transfer/UGATIT_100w/model.py new file mode 100644 index 00000000..f1639493 --- /dev/null +++ b/hub_module/modules/image/style_transfer/UGATIT_100w/model.py @@ -0,0 +1,67 @@ +import os +import numpy as np + +from paddle.fluid.core import AnalysisConfig, create_paddle_predictor + +__all__ = ['Model'] + +class Model(): + # 初始化函数 + def __init__(self, modelpath, use_gpu): + # 加载模型预测器 + self.predictor = self.load_model(modelpath, use_gpu) + + # 获取模型的输入输出 + self.input_names = self.predictor.get_input_names() + self.output_names = self.predictor.get_output_names() + self.input_tensor = self.predictor.get_input_tensor(self.input_names[0]) + self.output_tensor = self.predictor.get_output_tensor(self.output_names[0]) + + # 模型加载函数 + def load_model(self, modelpath, use_gpu): + # 对运行位置进行配置 + if use_gpu: + try: + places = os.environ["CUDA_VISIBLE_DEVICES"] + places = int(places[0]) + except Exception as e: + print('Error: %s. Please set the environment variables "CUDA_VISIBLE_DEVICES".' % e) + use_gpu = False + + # 加载模型参数 + config = AnalysisConfig(modelpath) + + # 设置参数 + if use_gpu: + config.enable_use_gpu(100, places) + else: + config.disable_gpu() + config.disable_glog_info() + config.switch_ir_optim(True) + config.enable_memory_optim() + config.switch_use_feed_fetch_ops(False) + config.switch_specify_input_names(True) + + # 通过参数加载模型预测器 + predictor = create_paddle_predictor(config) + + # 返回预测器 + return predictor + + # 模型预测函数 + def predict(self, input_datas): + outputs = [] + + # 遍历输入数据进行预测 + for input_data in input_datas: + inputs = input_data.copy() + self.input_tensor.copy_from_cpu(inputs) + self.predictor.zero_copy_run() + output = self.output_tensor.copy_to_cpu() + outputs.append(output) + + # 预测结果合并 + outputs = np.concatenate(outputs, 0) + + # 返回预测结果 + return outputs \ No newline at end of file diff --git a/hub_module/modules/image/style_transfer/UGATIT_100w/module.py b/hub_module/modules/image/style_transfer/UGATIT_100w/module.py new file mode 100644 index 00000000..5a2bdd0f --- /dev/null +++ b/hub_module/modules/image/style_transfer/UGATIT_100w/module.py @@ -0,0 +1,69 @@ +import os + +from paddlehub import Module +from paddlehub.module.module import moduleinfo, serving + +from UGATIT_100w.model import Model +from UGATIT_100w.processor import base64_to_cv2, cv2_to_base64, Processor + +@moduleinfo( + name="UGATIT_100w", # 模型名称 + type="CV/style_transfer", # 模型类型 + author="jm12138", # 作者名称 + author_email="jm12138@qq.com", # 作者邮箱 + summary="UGATIT_100w", # 模型介绍 + version="1.0.0" # 版本号 +) +class UGATIT_100w(Module): + # 初始化函数 + def _initialize(self, use_gpu=False): + # 设置模型路径 + self.model_path = os.path.join(self.directory, "UGATIT_100w") + + # 加载模型 + self.model = Model(self.model_path, use_gpu) + + # 关键点检测函数 + def style_transfer( + self, + images=None, + paths=None, + batch_size=1, + output_dir='output', + visualization=False + ): + # 加载数据处理器 + processor = Processor( + images, + paths, + output_dir, + batch_size + ) + + # 模型预测 + outputs = self.model.predict(processor.input_datas) + + + # 结果后处理 + results = processor.postprocess(outputs, visualization) + + # 返回结果 + return results + + # Hub Serving + @serving + def serving_method(self, images, **kwargs): + # 获取输入数据 + images_decode = [base64_to_cv2(image) for image in images] + + # 图片风格转换 + results = self.style_transfer(images_decode, **kwargs) + + # 对输出图片进行编码 + encodes = [] + for result in results: + encode = cv2_to_base64(result) + encodes.append(encode) + + # 返回结果 + return encodes \ No newline at end of file diff --git a/hub_module/modules/image/style_transfer/UGATIT_100w/processor.py b/hub_module/modules/image/style_transfer/UGATIT_100w/processor.py new file mode 100644 index 00000000..2589809d --- /dev/null +++ b/hub_module/modules/image/style_transfer/UGATIT_100w/processor.py @@ -0,0 +1,120 @@ +import os +import cv2 +import time +import base64 +import numpy as np + +__all__ = ['base64_to_cv2', 'cv2_to_base64', 'Processor'] + +def check_dir(dir_path): + # 目录检查函数 + if not os.path.exists(dir_path): + os.makedirs(dir_path) + elif os.path.isfile(dir_path): + os.remove(dir_path) + os.makedirs(dir_path) + +def base64_to_cv2(b64str): + # base64转cv2函数 + data = base64.b64decode(b64str.encode('utf8')) + data = np.frombuffer(data, np.uint8) + data = cv2.imdecode(data, cv2.IMREAD_COLOR) + return data + +def cv2_to_base64(image): + # cv2转base64函数 + data = cv2.imencode('.jpg', image)[1] + return base64.b64encode(data.tostring()).decode('utf8') + +class Processor(): + # 初始化函数 + def __init__( + self, + images=None, + paths=None, + output_dir='output', + batch_size=1 + ): + # 变量设置 + self.images = images + self.paths = paths + self.output_dir = output_dir + self.batch_size = batch_size + + # 获取原始输入数据 + self.datas = self.load_datas() + + # 对原始输入数据进行预处理 + self.input_datas = self.preprocess() + + # 读取数据函数 + def load_datas(self): + datas = [] + + # 读取数据列表 + if self.paths is not None: + for im_path in self.paths: + assert os.path.isfile(im_path), "The {} isn't a valid file path.".format(im_path) + im = cv2.imread(im_path) + datas.append(im) + + if self.images is not None: + datas = self.images + + # 返回数据列表 + return datas + + # 数据预处理函数 + def preprocess(self): + input_datas = [] + + # 数据预处理 + for i, img in enumerate(self.datas): + # 图像缩放 + img = cv2.resize(img, (256, 256)) + + # 归一化 + img = (img.astype('float32')/255.0 - 0.5) / 0.5 + + # 转置 + img = img.transpose((2, 0, 1)) + + # 增加维度 + img = np.expand_dims(img, axis=0) + + # 加入输入数据列表 + input_datas.append(img) + + # 数据按batch_size切分 + input_datas = np.concatenate(input_datas, 0) + split_num = len(self.datas)//self.batch_size+1 if len(self.datas)%self.batch_size!=0 else len(self.datas)//self.batch_size + input_datas = np.array_split(input_datas, split_num) + + # 返回预处理完成的数据 + return input_datas + + def postprocess(self, outputs, visualization): + results = [] + + for im_id, output in enumerate(outputs): + # 图像后处理 + img = (output * 0.5 + 0.5) * 255. + + # 限幅 + img = np.clip(img, 0, 255).astype(np.uint8) + + # 转置 + img = img.transpose((1, 2, 0)) + + # 可视化 + if visualization: + # 检查输出目录 + check_dir(self.output_dir) + + # 写入输出图片 + cv2.imwrite(os.path.join(self.output_dir, '%d_%d.jpg' % (im_id, time.time())), img) + + results.append(img) + + # 返回结果 + return results \ No newline at end of file diff --git a/hub_module/modules/image/style_transfer/UGATIT_83w/README.md b/hub_module/modules/image/style_transfer/UGATIT_83w/README.md new file mode 100644 index 00000000..19603baa --- /dev/null +++ b/hub_module/modules/image/style_transfer/UGATIT_83w/README.md @@ -0,0 +1,122 @@ +## 模型概述 +UGATIT 图像风格转换模型 + +模型可将输入的人脸图像转换成动漫风格 + +模型权重来自UGATIT-Paddle开源项目 + +模型所使用的权重为genA2B_0835000 + +模型详情请参考[UGATIT-Paddle开源项目](https://github.com/miraiwk/UGATIT-paddle) + +## 模型安装 + +```shell +$hub install UGATIT_83w +``` + + +## API 说明 + +```python +def style_transfer( + self, + images=None, + paths=None, + batch_size=1, + output_dir='output', + visualization=False +) +``` + +风格转换API，将输入的人脸图像转换成动漫风格。 + +转换效果图如下： + +![输入图像](https://ai-studio-static-online.cdn.bcebos.com/d130fabd8bd34e53b2f942b3766eb6bbd3c19c0676d04abfbd5cc4b83b66f8b6) +![输出图像](https://ai-studio-static-online.cdn.bcebos.com/78653331ee2d472b81ff5bbccd6a904a80d2c5208f9c42c789b4f09a1ef46332) + +**参数** + +* images (list\[numpy.ndarray\]): 图片数据，ndarray.shape 为 \[H, W, C\]，默认为 None； +* paths (list\[str\]): 图片的路径，默认为 None； +* batch\_size (int): batch 的大小，默认设为 1； +* visualization (bool): 是否将识别结果保存为图片文件，默认设为 False； +* output\_dir (str): 图片的保存路径，默认设为 output。 + + +**返回** + +* res (list\[numpy.ndarray\]): 输出图像数据，ndarray.shape 为 \[H, W, C\]。 + + +## 预测代码示例 + +```python +import cv2 +import paddlehub as hub + +# 模型加载 +# use_gpu：是否使用GPU进行预测 +model = hub.Module('UGATIT_83w', use_gpu=False) + +# 模型预测 +result = model.style_transfer(images=[cv2.imread('/PATH/TO/IMAGE')]) + +# or +# result = model.style_transfer(paths=['/PATH/TO/IMAGE']) +``` + +## 服务部署 + +PaddleHub Serving可以部署一个在线图像风格转换服务。 + +## 第一步：启动PaddleHub Serving + +运行启动命令： +```shell +$ hub serving start -m UGATIT_w83 +``` + +这样就完成了一个图像风格转换的在线服务API的部署，默认端口号为8866。 + +**NOTE:** 如使用GPU预测，则需要在启动服务之前，请设置CUDA\_VISIBLE\_DEVICES环境变量，否则不用设置。 + +## 第二步：发送预测请求 + +配置好服务端，以下数行代码即可实现发送预测请求，获取预测结果 + +```python +import requests +import json +import cv2 +import base64 + + +def cv2_to_base64(image): + data = cv2.imencode('.jpg', image)[1] + return base64.b64encode(data.tostring()).decode('utf8') + + +# 发送HTTP请求 +data = {'images':[cv2_to_base64(cv2.imread("/PATH/TO/IMAGE"))]} +headers = {"Content-type": "application/json"} +url = "http://127.0.0.1:8866/predict/UGATIT_w83" +r = requests.post(url=url, headers=headers, data=json.dumps(data)) + +# 打印预测结果 +print(r.json()["results"]) +``` + + +## 模型相关信息 + +### 模型代码 + +https://github.com/miraiwk/UGATIT-paddle + +### 依赖 + +paddlepaddle >= 1.8.0 + +paddlehub >= 1.8.0 \ No newline at end of file diff --git a/hub_module/modules/image/style_transfer/UGATIT_83w/model.py b/hub_module/modules/image/style_transfer/UGATIT_83w/model.py new file mode 100644 index 00000000..f1639493 --- /dev/null +++ b/hub_module/modules/image/style_transfer/UGATIT_83w/model.py @@ -0,0 +1,67 @@ +import os +import numpy as np + +from paddle.fluid.core import AnalysisConfig, create_paddle_predictor + +__all__ = ['Model'] + +class Model(): + # 初始化函数 + def __init__(self, modelpath, use_gpu): + # 加载模型预测器 + self.predictor = self.load_model(modelpath, use_gpu) + + # 获取模型的输入输出 + self.input_names = self.predictor.get_input_names() + self.output_names = self.predictor.get_output_names() + self.input_tensor = self.predictor.get_input_tensor(self.input_names[0]) + self.output_tensor = self.predictor.get_output_tensor(self.output_names[0]) + + # 模型加载函数 + def load_model(self, modelpath, use_gpu): + # 对运行位置进行配置 + if use_gpu: + try: + places = os.environ["CUDA_VISIBLE_DEVICES"] + places = int(places[0]) + except Exception as e: + print('Error: %s. Please set the environment variables "CUDA_VISIBLE_DEVICES".' % e) + use_gpu = False + + # 加载模型参数 + config = AnalysisConfig(modelpath) + + # 设置参数 + if use_gpu: + config.enable_use_gpu(100, places) + else: + config.disable_gpu() + config.disable_glog_info() + config.switch_ir_optim(True) + config.enable_memory_optim() + config.switch_use_feed_fetch_ops(False) + config.switch_specify_input_names(True) + + # 通过参数加载模型预测器 + predictor = create_paddle_predictor(config) + + # 返回预测器 + return predictor + + # 模型预测函数 + def predict(self, input_datas): + outputs = [] + + # 遍历输入数据进行预测 + for input_data in input_datas: + inputs = input_data.copy() + self.input_tensor.copy_from_cpu(inputs) + self.predictor.zero_copy_run() + output = self.output_tensor.copy_to_cpu() + outputs.append(output) + + # 预测结果合并 + outputs = np.concatenate(outputs, 0) + + # 返回预测结果 + return outputs \ No newline at end of file diff --git a/hub_module/modules/image/style_transfer/UGATIT_83w/module.py b/hub_module/modules/image/style_transfer/UGATIT_83w/module.py new file mode 100644 index 00000000..7baf4584 --- /dev/null +++ b/hub_module/modules/image/style_transfer/UGATIT_83w/module.py @@ -0,0 +1,69 @@ +import os + +from paddlehub import Module +from paddlehub.module.module import moduleinfo, serving + +from UGATIT_83w.model import Model +from UGATIT_83w.processor import base64_to_cv2, cv2_to_base64, Processor + +@moduleinfo( + name="UGATIT_83w", # 模型名称 + type="CV/style_transfer", # 模型类型 + author="jm12138", # 作者名称 + author_email="jm12138@qq.com", # 作者邮箱 + summary="UGATIT", # 模型介绍 + version="1.0.0" # 版本号 +) +class UGATIT_83w(Module): + # 初始化函数 + def _initialize(self, use_gpu=False): + # 设置模型路径 + self.model_path = os.path.join(self.directory, "UGATIT_83w") + + # 加载模型 + self.model = Model(self.model_path, use_gpu) + + # 关键点检测函数 + def style_transfer( + self, + images=None, + paths=None, + batch_size=1, + output_dir='output', + visualization=False + ): + # 加载数据处理器 + processor = Processor( + images, + paths, + output_dir, + batch_size + ) + + # 模型预测 + outputs = self.model.predict(processor.input_datas) + + + # 结果后处理 + results = processor.postprocess(outputs, visualization) + + # 返回结果 + return results + + # Hub Serving + @serving + def serving_method(self, images, **kwargs): + # 获取输入数据 + images_decode = [base64_to_cv2(image) for image in images] + + # 图片风格转换 + results = self.style_transfer(images_decode, **kwargs) + + # 对输出图片进行编码 + encodes = [] + for result in results: + encode = cv2_to_base64(result) + encodes.append(encode) + + # 返回结果 + return encodes \ No newline at end of file diff --git a/hub_module/modules/image/style_transfer/UGATIT_83w/processor.py b/hub_module/modules/image/style_transfer/UGATIT_83w/processor.py new file mode 100644 index 00000000..2589809d --- /dev/null +++ b/hub_module/modules/image/style_transfer/UGATIT_83w/processor.py @@ -0,0 +1,120 @@ +import os +import cv2 +import time +import base64 +import numpy as np + +__all__ = ['base64_to_cv2', 'cv2_to_base64', 'Processor'] + +def check_dir(dir_path): + # 目录检查函数 + if not os.path.exists(dir_path): + os.makedirs(dir_path) + elif os.path.isfile(dir_path): + os.remove(dir_path) + os.makedirs(dir_path) + +def base64_to_cv2(b64str): + # base64转cv2函数 + data = base64.b64decode(b64str.encode('utf8')) + data = np.frombuffer(data, np.uint8) + data = cv2.imdecode(data, cv2.IMREAD_COLOR) + return data + +def cv2_to_base64(image): + # cv2转base64函数 + data = cv2.imencode('.jpg', image)[1] + return base64.b64encode(data.tostring()).decode('utf8') + +class Processor(): + # 初始化函数 + def __init__( + self, + images=None, + paths=None, + output_dir='output', + batch_size=1 + ): + # 变量设置 + self.images = images + self.paths = paths + self.output_dir = output_dir + self.batch_size = batch_size + + # 获取原始输入数据 + self.datas = self.load_datas() + + # 对原始输入数据进行预处理 + self.input_datas = self.preprocess() + + # 读取数据函数 + def load_datas(self): + datas = [] + + # 读取数据列表 + if self.paths is not None: + for im_path in self.paths: + assert os.path.isfile(im_path), "The {} isn't a valid file path.".format(im_path) + im = cv2.imread(im_path) + datas.append(im) + + if self.images is not None: + datas = self.images + + # 返回数据列表 + return datas + + # 数据预处理函数 + def preprocess(self): + input_datas = [] + + # 数据预处理 + for i, img in enumerate(self.datas): + # 图像缩放 + img = cv2.resize(img, (256, 256)) + + # 归一化 + img = (img.astype('float32')/255.0 - 0.5) / 0.5 + + # 转置 + img = img.transpose((2, 0, 1)) + + # 增加维度 + img = np.expand_dims(img, axis=0) + + # 加入输入数据列表 + input_datas.append(img) + + # 数据按batch_size切分 + input_datas = np.concatenate(input_datas, 0) + split_num = len(self.datas)//self.batch_size+1 if len(self.datas)%self.batch_size!=0 else len(self.datas)//self.batch_size + input_datas = np.array_split(input_datas, split_num) + + # 返回预处理完成的数据 + return input_datas + + def postprocess(self, outputs, visualization): + results = [] + + for im_id, output in enumerate(outputs): + # 图像后处理 + img = (output * 0.5 + 0.5) * 255. + + # 限幅 + img = np.clip(img, 0, 255).astype(np.uint8) + + # 转置 + img = img.transpose((1, 2, 0)) + + # 可视化 + if visualization: + # 检查输出目录 + check_dir(self.output_dir) + + # 写入输出图片 + cv2.imwrite(os.path.join(self.output_dir, '%d_%d.jpg' % (im_id, time.time())), img) + + results.append(img) + + # 返回结果 + return results \ No newline at end of file diff --git a/hub_module/modules/image/style_transfer/UGATIT_92w/README.md b/hub_module/modules/image/style_transfer/UGATIT_92w/README.md new file mode 100644 index 00000000..b6edd1ff --- /dev/null +++ b/hub_module/modules/image/style_transfer/UGATIT_92w/README.md @@ -0,0 +1,122 @@ +## 模型概述 +UGATIT 图像风格转换模型 + +模型可将输入的人脸图像转换成动漫风格 + +模型权重来自UGATIT-Paddle开源项目 + +模型所使用的权重为genA2B_0924000 + +模型详情请参考[UGATIT-Paddle开源项目](https://github.com/miraiwk/UGATIT-paddle) + +## 模型安装 + +```shell +$hub install UGATIT_92w +``` + + +## API 说明 + +```python +def style_transfer( + self, + images=None, + paths=None, + batch_size=1, + output_dir='output', + visualization=False +) +``` + +风格转换API，将输入的人脸图像转换成动漫风格。 + +转换效果图如下： + +![输入图像](https://ai-studio-static-online.cdn.bcebos.com/d130fabd8bd34e53b2f942b3766eb6bbd3c19c0676d04abfbd5cc4b83b66f8b6) +![输出图像](https://ai-studio-static-online.cdn.bcebos.com/b7305162ff6345e9b04507a196ebe854907b446936934844be8aae4b0297db18) + +**参数** + +* images (list\[numpy.ndarray\]): 图片数据，ndarray.shape 为 \[H, W, C\]，默认为 None； +* paths (list\[str\]): 图片的路径，默认为 None； +* batch\_size (int): batch 的大小，默认设为 1； +* visualization (bool): 是否将识别结果保存为图片文件，默认设为 False； +* output\_dir (str): 图片的保存路径，默认设为 output。 + + +**返回** + +* res (list\[numpy.ndarray\]): 输出图像数据，ndarray.shape 为 \[H, W, C\]。 + + +## 预测代码示例 + +```python +import cv2 +import paddlehub as hub + +# 模型加载 +# use_gpu：是否使用GPU进行预测 +model = hub.Module('UGATIT_92w', use_gpu=False) + +# 模型预测 +result = model.style_transfer(images=[cv2.imread('/PATH/TO/IMAGE')]) + +# or +# result = model.style_transfer(paths=['/PATH/TO/IMAGE']) +``` + +## 服务部署 + +PaddleHub Serving可以部署一个在线图像风格转换服务。 + +## 第一步：启动PaddleHub Serving + +运行启动命令： +```shell +$ hub serving start -m UGATIT_92w +``` + +这样就完成了一个图像风格转换的在线服务API的部署，默认端口号为8866。 + +**NOTE:** 如使用GPU预测，则需要在启动服务之前，请设置CUDA\_VISIBLE\_DEVICES环境变量，否则不用设置。 + +## 第二步：发送预测请求 + +配置好服务端，以下数行代码即可实现发送预测请求，获取预测结果 + +```python +import requests +import json +import cv2 +import base64 + + +def cv2_to_base64(image): + data = cv2.imencode('.jpg', image)[1] + return base64.b64encode(data.tostring()).decode('utf8') + + +# 发送HTTP请求 +data = {'images':[cv2_to_base64(cv2.imread("/PATH/TO/IMAGE"))]} +headers = {"Content-type": "application/json"} +url = "http://127.0.0.1:8866/predict/UGATIT_92w" +r = requests.post(url=url, headers=headers, data=json.dumps(data)) + +# 打印预测结果 +print(r.json()["results"]) +``` + + +## 模型相关信息 + +### 模型代码 + +https://github.com/miraiwk/UGATIT-paddle + +### 依赖 + +paddlepaddle >= 1.8.0 + +paddlehub >= 1.8.0 \ No newline at end of file diff --git a/hub_module/modules/image/style_transfer/UGATIT_92w/model.py b/hub_module/modules/image/style_transfer/UGATIT_92w/model.py new file mode 100644 index 00000000..f1639493 --- /dev/null +++ b/hub_module/modules/image/style_transfer/UGATIT_92w/model.py @@ -0,0 +1,67 @@ +import os +import numpy as np + +from paddle.fluid.core import AnalysisConfig, create_paddle_predictor + +__all__ = ['Model'] + +class Model(): + # 初始化函数 + def __init__(self, modelpath, use_gpu): + # 加载模型预测器 + self.predictor = self.load_model(modelpath, use_gpu) + + # 获取模型的输入输出 + self.input_names = self.predictor.get_input_names() + self.output_names = self.predictor.get_output_names() + self.input_tensor = self.predictor.get_input_tensor(self.input_names[0]) + self.output_tensor = self.predictor.get_output_tensor(self.output_names[0]) + + # 模型加载函数 + def load_model(self, modelpath, use_gpu): + # 对运行位置进行配置 + if use_gpu: + try: + places = os.environ["CUDA_VISIBLE_DEVICES"] + places = int(places[0]) + except Exception as e: + print('Error: %s. Please set the environment variables "CUDA_VISIBLE_DEVICES".' % e) + use_gpu = False + + # 加载模型参数 + config = AnalysisConfig(modelpath) + + # 设置参数 + if use_gpu: + config.enable_use_gpu(100, places) + else: + config.disable_gpu() + config.disable_glog_info() + config.switch_ir_optim(True) + config.enable_memory_optim() + config.switch_use_feed_fetch_ops(False) + config.switch_specify_input_names(True) + + # 通过参数加载模型预测器 + predictor = create_paddle_predictor(config) + + # 返回预测器 + return predictor + + # 模型预测函数 + def predict(self, input_datas): + outputs = [] + + # 遍历输入数据进行预测 + for input_data in input_datas: + inputs = input_data.copy() + self.input_tensor.copy_from_cpu(inputs) + self.predictor.zero_copy_run() + output = self.output_tensor.copy_to_cpu() + outputs.append(output) + + # 预测结果合并 + outputs = np.concatenate(outputs, 0) + + # 返回预测结果 + return outputs \ No newline at end of file diff --git a/hub_module/modules/image/style_transfer/UGATIT_92w/module.py b/hub_module/modules/image/style_transfer/UGATIT_92w/module.py new file mode 100644 index 00000000..b6a8723a --- /dev/null +++ b/hub_module/modules/image/style_transfer/UGATIT_92w/module.py @@ -0,0 +1,69 @@ +import os + +from paddlehub import Module +from paddlehub.module.module import moduleinfo, serving + +from UGATIT_92w.model import Model +from UGATIT_92w.processor import base64_to_cv2, cv2_to_base64, Processor + +@moduleinfo( + name="UGATIT_92w", # 模型名称 + type="CV/style_transfer", # 模型类型 + author="jm12138", # 作者名称 + author_email="jm12138@qq.com", # 作者邮箱 + summary="UGATIT_92w", # 模型介绍 + version="1.0.0" # 版本号 +) +class UGATIT_92w(Module): + # 初始化函数 + def _initialize(self, use_gpu=False): + # 设置模型路径 + self.model_path = os.path.join(self.directory, "UGATIT_92w") + + # 加载模型 + self.model = Model(self.model_path, use_gpu) + + # 关键点检测函数 + def style_transfer( + self, + images=None, + paths=None, + batch_size=1, + output_dir='output', + visualization=False + ): + # 加载数据处理器 + processor = Processor( + images, + paths, + output_dir, + batch_size + ) + + # 模型预测 + outputs = self.model.predict(processor.input_datas) + + + # 结果后处理 + results = processor.postprocess(outputs, visualization) + + # 返回结果 + return results + + # Hub Serving + @serving + def serving_method(self, images, **kwargs): + # 获取输入数据 + images_decode = [base64_to_cv2(image) for image in images] + + # 图片风格转换 + results = self.style_transfer(images_decode, **kwargs) + + # 对输出图片进行编码 + encodes = [] + for result in results: + encode = cv2_to_base64(result) + encodes.append(encode) + + # 返回结果 + return encodes \ No newline at end of file diff --git a/hub_module/modules/image/style_transfer/UGATIT_92w/processor.py b/hub_module/modules/image/style_transfer/UGATIT_92w/processor.py new file mode 100644 index 00000000..2589809d --- /dev/null +++ b/hub_module/modules/image/style_transfer/UGATIT_92w/processor.py @@ -0,0 +1,120 @@ +import os +import cv2 +import time +import base64 +import numpy as np + +__all__ = ['base64_to_cv2', 'cv2_to_base64', 'Processor'] + +def check_dir(dir_path): + # 目录检查函数 + if not os.path.exists(dir_path): + os.makedirs(dir_path) + elif os.path.isfile(dir_path): + os.remove(dir_path) + os.makedirs(dir_path) + +def base64_to_cv2(b64str): + # base64转cv2函数 + data = base64.b64decode(b64str.encode('utf8')) + data = np.frombuffer(data, np.uint8) + data = cv2.imdecode(data, cv2.IMREAD_COLOR) + return data + +def cv2_to_base64(image): + # cv2转base64函数 + data = cv2.imencode('.jpg', image)[1] + return base64.b64encode(data.tostring()).decode('utf8') + +class Processor(): + # 初始化函数 + def __init__( + self, + images=None, + paths=None, + output_dir='output', + batch_size=1 + ): + # 变量设置 + self.images = images + self.paths = paths + self.output_dir = output_dir + self.batch_size = batch_size + + # 获取原始输入数据 + self.datas = self.load_datas() + + # 对原始输入数据进行预处理 + self.input_datas = self.preprocess() + + # 读取数据函数 + def load_datas(self): + datas = [] + + # 读取数据列表 + if self.paths is not None: + for im_path in self.paths: + assert os.path.isfile(im_path), "The {} isn't a valid file path.".format(im_path) + im = cv2.imread(im_path) + datas.append(im) + + if self.images is not None: + datas = self.images + + # 返回数据列表 + return datas + + # 数据预处理函数 + def preprocess(self): + input_datas = [] + + # 数据预处理 + for i, img in enumerate(self.datas): + # 图像缩放 + img = cv2.resize(img, (256, 256)) + + # 归一化 + img = (img.astype('float32')/255.0 - 0.5) / 0.5 + + # 转置 + img = img.transpose((2, 0, 1)) + + # 增加维度 + img = np.expand_dims(img, axis=0) + + # 加入输入数据列表 + input_datas.append(img) + + # 数据按batch_size切分 + input_datas = np.concatenate(input_datas, 0) + split_num = len(self.datas)//self.batch_size+1 if len(self.datas)%self.batch_size!=0 else len(self.datas)//self.batch_size + input_datas = np.array_split(input_datas, split_num) + + # 返回预处理完成的数据 + return input_datas + + def postprocess(self, outputs, visualization): + results = [] + + for im_id, output in enumerate(outputs): + # 图像后处理 + img = (output * 0.5 + 0.5) * 255. + + # 限幅 + img = np.clip(img, 0, 255).astype(np.uint8) + + # 转置 + img = img.transpose((1, 2, 0)) + + # 可视化 + if visualization: + # 检查输出目录 + check_dir(self.output_dir) + + # 写入输出图片 + cv2.imwrite(os.path.join(self.output_dir, '%d_%d.jpg' % (im_id, time.time())), img) + + results.append(img) + + # 返回结果 + return results \ No newline at end of file -- GitLab