AI画家第四弹——利用Flask发布风格迁移API

上篇文章介绍了python web开发中常用到的一个框架flask，若是有遗忘的，能够点此回顾👉AI画家第三弹——毕业设计大杀器之Flask，本文的主要任务就是完成上篇文章末尾的要求，利用Flask发布你本身的风格迁移API。

本文源码可在微信公众号「01二进制」后台回复「风格迁移API」得到

需求分析

咱们知道软件工程的第一步就是需求分析，放在这里就是要知道咱们须要实现的功能是什么样的。我画了一张简陋的图来描述此次的需求：

真的是很简陋的一张图啊，其实理解起来很容易，就是用户上传一张图片，Flask获取到这张图片，调用风格迁移的模型，而后生成结果图，在传递回前端便可。

环境准备

既然明白了需求，那么接下来要作的事情天然是环境搭建了，老样子，这里咱们仍然使用Pipenv来建立虚拟环境，如何搭建pipenv环境我就不说了，在微信公众号「01二进制」后台回复「风格迁移API」得到源码以后直接在终端输入pipenv install便可。

开始

hello world

咱们首先在项目根目录建立一个main.py的文件做为整个项目的启动文件，上文咱们说过，为了简化大型应用并为扩展提供集中的注册入口，咱们并不会将全部的视图函数直接写在main.py，而是采用蓝图的方式分模块开发，所以咱们须要在项目根目录新建app/文件夹，在其中的__init__.py中编写以下代码：

from flask import Flask

def create_app():
    app = Flask(__name__)
    return app
复制代码

这样咱们就能够经过在main.py中编写以下代码实现一个hello world应用了。

from app import create_app

app = create_app()

@app.route('/')
def hello():
    return 'hello,world'

if __name__ == '__main__':
    app.run(port=8080, debug=True)
复制代码

启动main.py便可发现项目启动了，在浏览器输入localhost:8080便可看到hello,world字样。

蓝图编写

咱们确定是不能知足于小小的hello world的，既然说到了模块化开发，那怎么个模块法？

这里每一个人的想法都是不同的，其实也没有一个统一的标准，这里我就说下我本身的分级方法吧。

• stylize是项目根路径
• app是项目
• app/api是项目的api部分，一个项目确定不仅有api，还可能会有web等页面内容
• app/api/v1代表该api的版本是v1，固然往后也有可能会有v二、v3等等
• app/api/v1/img代表这里存放的都是和img有关的api
• app/api/v1/img/stylize.py代表这个文件存放的是风格迁移的视图函数

认识完结构的划分以后，就来编写咱们的蓝图吧。

首先咱们须要在app/api/v1/img/__init__.py中编写以下代码：

from flask import Blueprint

# 定义一个蓝图
img = Blueprint('img', __name__)

from app.api.v1.img import stylize

# 这段代码用来测试该接口是否可用
@img.route('/')
def say_hello():
    return '这里是图片处理类的接口'
复制代码

这样咱们就定义了一个叫作img的蓝图，而后咱们在app/api/v1/__init__.py中编写以下代码：

from flask import Blueprint

# 定义一个蓝图
v1 = Blueprint('v1', __name__)

from app.api.v1.img import img
复制代码

这样咱们就实现了v1蓝图的编写。

那这样是否是就可使用蓝图了呢？固然不是，咱们还须要在app中配置这个蓝图，把蓝图加载到app中，不然flask是没法识别蓝图的。加载的方法也很简单，咱们在app/__init__.py文件中添加一个函数：

def register_blueprint(app):
    from app.api.v1 import v1
    from app.api.v1.img import img

    app.register_blueprint(v1, url_prefix='/api/v1')
    app.register_blueprint(img, url_prefix='/api/v1/img')
复制代码

将这两个蓝图注册到app中，其中url_prefix这个参数用来标注路由的。有人可能不清楚，这里举个例子你就懂了。咱们启动这个项目以后，在浏览器输入的是localhost:8080，若是加了url_prefix这个参数以后，咱们访问img下的视图函数时就须要把路径改成localhost:8080/api/v1/img/了。

而后在create_app()函数中调用这个方法便可，main.py的代码以下：

from flask import Flask


def create_app():
    app = Flask(__name__)

    register_blueprint(app)
    return app


def register_blueprint(app):
    from app.api.v1 import v1
    from app.api.v1.img import img

    app.register_blueprint(v1, url_prefix='/api/v1')
    app.register_blueprint(img, url_prefix='/api/v1/img')


if __name__ == '__main__':
    create_app()
复制代码

接下来咱们测试下这个蓝图是否真的注册成功了，咱们启动该项目，并打开Postman输入：localhost:8080/api/v1/img，咱们能够看到以下信息：

说明咱们的蓝图已经注册

编写风格迁移工具类

既然蓝图都已经编写好了，那么视图函数的编写也就很是简单了，所以这里咱们先把风格迁移的工具类写好，最后再编写视图函数。

上上篇文章咱们介绍了图像风格迁移，记不清的能够看这篇文章👉AI绘画第二弹——图像风格迁移，这篇文章介绍了最传统的图像风格迁移，想要生成一张图片的速度很是很是慢，确定是没办法做为实际使用的，所以这篇文章所采用的生成风格迁移的图片的方法并非这篇文章，而是基于李飞飞等人的一篇论文《Perceptual Losses for Real-Time Style Transfer and Super-Resolution》所实现的快速图像风格迁移。这里只介绍如何拿训练好的模型去运用，有兴趣的本身下载这篇文章去研究。

咱们新建一个文件夹：app/utils/stylize，这个文件夹中包含了风格迁移的工具类，项目结构以下：

其中output为最终的生成文件夹、src存放风格迁移的文件（能够不用管）、evaluate.py是以前用于评估模型性能的文件（其实也就是生成图片的文件，不用管+1）、trained_model文件夹存放了咱们已经训练好的模型、temp存放的是咱们从前端上传的图片，create_stylize_photo.py包含的是咱们对外提供的风格迁移工具类。

所以这整个文件夹咱们只须要关注create_stylize_photo.py这一个文件就能够了，其余的下载源码以后本身看就能够了。

压缩图片

# 压缩图片
def compress_image():
    im = Image.open(content_image)
    if content_image.endswith(".png"):
        im = im.convert('P')
    im.save(content_image, optimize=True)
复制代码

选择图片风格即所须要使用的模型

style_list = ['la_muse', 'rain_princess', 'scream', 'udnie', 'wave', 'wreck']
style = style_list[int(image_style)]

# 模型的 checkpoint 的位置
check_point_dir = trained_models_path + style + '.ckpt'
复制代码

执行生成图片的操做

# 最终生成的图片路径
result_image = path + '/output/' + 'output.jpg'

# 执行生成图片的操做
ffwd_to_img(content_image, result_image, check_point_dir)
复制代码

执行完上述步骤后，咱们就能够在utils/stylize/output中看到已经生成的风格迁移图片了。

利用七牛云存储结果图片

因为服务器带宽限制，咱们最后返回给前端的结果确定不能是咱们本身服务器的url（毕竟学生机的带宽只有1M），因此这里我建议使用七牛云存储的功能将生成的结果保存到七牛云上，而后返回一个url便可。

# 将生成的图片上传到七牛云
# 传入filename和filepath，返回图片的URL
def upload_pic_to_qiniu(filename, filepath):
    from app.secure import QINIU_AK
    from app.secure import QINIU_SK

    access_key = QINIU_AK
    secret_key = QINIU_SK

    q = Auth(access_key, secret_key)

    # 要上传的空间
    bucket_name = 'ytools'

    # 生成上传 Token，能够指定过时时间等
    token = q.upload_token(bucket_name, filename, 3600)

    ret, info = put_file(token, filename, filepath)
    return BASE_URL + ret['key']
复制代码

而后咱们在create_stylize_photo.py中加一句存储的代码便可。

img_url = upload_pic_to_qiniu(filename, result_image)
复制代码

编写风格迁移API

如今咱们经过风格迁移工具类已经能够实现输入一张原始图片返回生成图片的URL的功能，如今咱们来将目光聚焦到风格迁移API的编写上。

定义路由

@img.route('/stylize/create', methods=['POST'])
def create_style_changed_img():
复制代码

接收参数

咱们定义一个函数create_style_changed_img()，方法采用POST方式，咱们须要接受前端发来的两个参数，分别是img和type

img = request.files.get('img')
type = request.form.get('type')
复制代码

生成图片及URL

而后咱们将接收到的文件保存到以前新建的temp文件加中，而后调用工具类的方法返回图片的url

img.save(path + '/temp/' + 'temp.jpg')

img_url = change_style(int(type))
复制代码

定义返回格式

做为一个好的api，咱们确定不仅能返回一张图片的url就能够了，咱们还须要记录下生成的时间，所以咱们在代码执行的开始和结束的时候分别添加一段代码：

start = datetime.datetime.now()
end = datetime.datetime.now()
复制代码

而后咱们再定义返回格式

status = 200
msg = '图片生成成功'
info = [
    {
        'img_url': img_url,
        'created_time': get_date_now(),
        'finish_time': (end - start).seconds
    }
]
复制代码

而后将结果返回

res_json = Res(status, msg, info)

return jsonify(res_json.__dict__)
复制代码

这里的Res是我定义的一个返回的实体信息类，长得是下面这样

class Res:
    status = 200
    msg = ''
    info = []

    def __init__(self, status, msg, info):
        self.status = status
        self.msg = msg
        self.info = info
复制代码

最后咱们调用flask的jsonify方法，就能够返回json结果。

测试

到这里咱们就已经完成了一个风格迁移API的编写，接下来咱们测试下咱们的API吧，首先先启动项目，而后打开Postman，将请求方法改成post，添加两个参数img和type，以下：

选择图片的时候，图片的质量尽可能不要太大，不然可能会出现卡死的状况

最后的返回结果以下：

{
    "info": [
        {
            "created_time": "2019-05-15 15:17:25",
            "finish_time": 2,
            "img_url": "https://user-gold-cdn.xitu.io/2019/5/15/16aba682a31521a7?w=640&h=360&f=jpeg&s=39700"
        }
    ],
    "msg": "图片生成成功",
    "status": 200
}
复制代码

访问该URL便可看到以下图片（感受还蛮好看的）：

至此咱们已经实现了利用Flask发布一个风格迁移API了。

总结

本文源码可在微信公众号「01二进制」后台回复「风格迁移API」得到

最后总结下，在这篇文章中介绍了如何利用Flask发布一个风格迁移API，其中咱们介绍了应该如何利用蓝图进行模块化开发，并给出了我本身认为的比较好的分层方法，同时利用七牛云存储为咱们的服务器减压，最后利用postman请求该API完成测试。下一篇将会介绍如何将该API部署到腾讯云（阿里云）上实现外网访问，有兴趣的小伙伴给个关注，你的支持才是我更新的最大动力！