基于飞桨PaddlePaddle实现素描线稿提取——仅需一行代码即可实现模型调用
一行代码实现从彩色图提取素描线稿
PaddleHub是飞桨预训练模型应用工具,能拿来即用,一键预测。本项目以开发者七年期限的项目——PaddlePaddle之视频转线稿篇为基础,将该模型贡献到PaddleHub的模型库里,并用一行代码实现快速加载。
一、前言:初识PaddleHub
PaddleHub官方Github地址(欢迎star~):
https://github.com/PaddlePaddle/PaddleHub
PaddleHub旨在为开发人员提供丰富,高质量,可直接使用的预训练模型。开发者无需深度学习背景,便可以快速使用PaddleHub内置的AI模型。其模型涵盖图像,文本,音频和视频的4个主要类别,并支持一键式预测,轻松的服务部署和迁移学习,另外,所有模型都是开源的,可以在离线情况下免费下载和使用。
预训练模型(Pre-Trained Models)
预训练模型是一个已经训练好的保存下来的网络。该网络之前已经在一个大型的数据集上进行了训练(如图像相关类的网络在ImageNet数据上进行了训练)
二、构建Module的大框架
PaddleHub Module是使用PaddleHub的基础。其可以通过指定名称即可方便地完成一键加载,如加载预训练模型ERNIE仅需一行代码即可完成,hub.Module(name=‘ernie’),省去了复杂的网络结构代码以及参数加载的繁琐过程。
将模型贡献到PaddleHub的模型库之前,我们必须搭好这一框架。
1.将模型文件存放在infer_model下
# 创建必要的文件夹
!mkdir -p Extract_Line_Draft/assets/infer_model
# 将准备好的模型文件解压
!unzip work/inference_model.zip
# 将模型文件复制到Module指定目录下
!cp -r inference_model/* Extract_Line_Draft/assets/infer_model
2.创建必要的文件__init__.py
init = open("Extract_Line_Draft/__init__.py", "w")
KaTeX parse error: Expected group after '_' at position 1: _̲_init__.py是一个空文件,直接创建即可。
3.创建必要的文件module.py
module = open("Extract_Line_Draft/module.py", "w")
导入必要的包
import argparse
import ast
import os
import math
import six
from pathlib import Path
from paddle.fluid.core import PaddleTensor, AnalysisConfig, create_paddle_predictor
from paddlehub.module.module import runnable, serving, moduleinfo
from paddlehub.io.parser import txt_parser
import numpy as np
import paddle.fluid as fluid
import paddlehub as hub
from Extract_Line_Draft.function import *
填写module的基本信息
@moduleinfo(
name="Extract_Line_Draft",
version="1.0.0",
type="cv/segmentation",
summary=
"Import the color picture and generate the line draft of the picture",
author="彭兆帅,郑博培",
author_email="1084667371@qq.com,2733821739@qq.com")
实现预测功能
class ExtractLineDraft(hub.Module):
def _initialize(self):
"""
Initialize with the necessary elements
"""
# 加载模型路径
self.pretrained_model_path = os.path.join(self.directory, "assets","infer_model")
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = AnalysisConfig(model, params)
else:
config = AnalysisConfig(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
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
def ExtractLine(self, image_path, use_gpu=False):
"""
Get the input and program of the infer model
Args:
image_path(str): Path to the folder containing the input picture
use_gpu(bool): Weather to use gpu
"""
for home, dirs, files in os.walk(image_path):
for filename in files:
fullname = os.path.join(home, filename)
from_mat = cv2.imread(fullname)
width = float(from_mat.shape[1])
height = float(from_mat.shape[0])
new_width = 0
new_height = 0
if (width > height):
from_mat = cv2.resize(from_mat, (512, int(512 / width * height)), interpolation=cv2.INTER_AREA)
new_width = 512
new_height = int(512 / width * height)
else:
from_mat = cv2.resize(from_mat, (int(512 / height * width), 512), interpolation=cv2.INTER_AREA)
new_width = int(512 / height * width)
new_height = 512
from_mat = from_mat.transpose((2, 0, 1))
light_map = np.zeros(from_mat.shape, dtype=np.float)
for channel in range(3):
light_map[channel] = get_light_map_single(from_mat[channel])
light_map = normalize_pic(light_map)
light_map = resize_img_512_3d(light_map)
light_map = light_map.astype('float32')
# 加载模型预测器
self.predictor = self.load_model(self.pretrained_model_path, use_gpu, use_mkldnn=True, combined=False)
# 获取模型的输入输出
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])
line_mat = self.predict(np.expand_dims(light_map, axis=0).astype('float32'))
# 去除 batch 维度 (512, 512, 3)
line_mat = line_mat.transpose((3, 1, 2, 0))[0]
# 裁剪 (512, 384, 3)
line_mat = line_mat[0:int(new_height), 0:int(new_width), :]
line_mat = np.amax(line_mat, 2)
# 保存图片
if Path('./output/').exists():
show_active_img_and_save_denoise(line_mat, './output/' + filename)
else:
os.makedirs('./output/')
show_active_img_and_save_denoise(line_mat, './output/' + filename)
print('图片' + filename + '已经完成')
print('全部图片转换成功。')
实现命令行调用的接口
@runnable
def run_cmd(self, argvs):
"""
Run as a command.
"""
self.parser = argparse.ArgumentParser(
description='Run the %s module.' % self.name,
prog='hub run %s' % self.name,
usage='%(prog)s',
add_help=True)
self.arg_input_group = self.parser.add_argument_group(
title="Input options", description="Input data. Required")
self.arg_config_group = self.parser.add_argument_group(
title="Config options",
description=
"Run configuration for controlling module behavior, not required.")
self.add_module_input_arg()
args = self.parser.parse_args(argvs)
try:
input_data = self.check_input_data(args)
except RuntimeError:
self.parser.print_help()
return None
use_gpu = args.use_gpu
self.ExtractLine(image_path=input_data, use_gpu=use_gpu)
def add_module_input_arg(self):
"""
Add the command input options
"""
self.arg_input_group.add_argument(
'--input_path',
type=str,
default=None,
help="file contain input data")
self.arg_input_group.add_argument(
'--use_gpu',
type=str,
default=None,
help="weather to use gpu")
def check_input_data(self, args):
input_data = []
if args.input_path:
if not os.path.exists(args.input_path):
raise RuntimeError("Path %s is not exist." % args.input_path)
path = "{}".format(args.input_path)
# print(path)
return path
三、测试及调用
1.安装并查看Module基本信息
# 安装Module
!hub install Extract_Line_Draft
# 查看该Module的基本信息
!hub show Extract_Line_Draft
2.通过API的方式调用
import paddlehub as hub
Extract_Line_Draft_test = hub.Module(name="Extract_Line_Draft")
test_img_path = "testImage"
# execute predict and print the result
t the result
Extract_Line_Draft_test.ExtractLine(test_img_path, use_gpu=True)
[32m[2020-12-08 02:29:43,474] [ INFO] - Installing Extract_Line_Draft module[0m
[32m[2020-12-08 02:29:43,524] [ INFO] - Module Extract_Line_Draft already installed in /home/aistudio/.paddlehub/modules/Extract_Line_Draft[0m
图片2.jpg已经完成
图片1.jpg已经完成
图片3.jpg已经完成
全部图片转换成功。
3.通过命令行的方式调用——一行代码实现从彩色图提取素描线稿
!hub run Extract_Line_Draft --input_path "testImage" --use_gpu True
图片2.jpg已经完成
图片1.jpg已经完成
图片3.jpg已经完成
全部图片转换成功。
4.效果展示
原图
线稿图
四、个人介绍
北京联合大学机器人学院自动化专业2018级本科生 郑博培
百度飞桨开发者技术专家 PPDE
深圳柴火创客空间 认证会员
百度大脑 智能对话训练师
来AI Studio互粉吧,等你哦~ https://aistudio.baidu.com/aistudio/personalcenter/thirdview/147378
本文地址:https://blog.csdn.net/zbp_12138/article/details/110849163