GFPGAN 集成Flask 接口化改造

news2024/12/26 22:30:48

GFPGAN是一款腾讯开源的人脸高清修复模型，基于github上提供的demo，可以简单的集成Flask以实现功能接口化。
GFPGAN的安装，Flask的安装请参见其他文章。
测试方式
如若使用POSTMAN进行测试，需使用POST方式，form-data的请求体，发送图片文件到服务端，服务端会直接返回图片流，可将响应体直接保存为文件。同样的，也可供给网页form表单使用。

import io
from flask import Flask, request, Response
import cv2
import numpy as np
import os

from gfpgan import GFPGANer

app = Flask(__name__)

# 定义修复者
restorer = GFPGANer(
    model_path=os.path.join('experiments/pretrained_models', 'GFPGANv1.3.pth'),
)


@app.route('/action', methods=['POST', 'GET'])
def action():
    file = request.files['file']
    img_name = file.filename
    input_img = cv2.imdecode(np.asarray(bytearray(file.read()), dtype=np.uint8), -1)
    basename, ext = os.path.splitext(img_name)

    _, _, restored_img = restorer.enhance(input_img)

    img_ret = cv2.imencode(f'.{ext}', restored_img)[1].tobytes()
    img_ret = io.BytesIO(img_ret)

    res = Response(img_ret.read())
    res.headers.add('Content-Type', 'image/' + ext)
    res.headers.add('Content-Disposition', f'attachment; filename=restore_{img_name}')
    return res


def main():
    app.run(port=2020, host="127.0.0.1", debug=True)


if __name__ == '__main__':
    main()