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部署到騰訊雲（阿里雲）上實現外網訪問，有興趣的小夥伴給個關注，你的支持纔是我更新的最大動力！