DeOldify 接口化改造集成 Flask

news2024/9/30 21:19:07

类似的图片修复项目 GFPGAN 的改造见我另一篇文
https://blog.csdn.net/weixin_43074462/article/details/132497146

DeOldify 是一款开源软件，用于给黑白照片或视频上色，效果还不错。

安装部署教程请参考别的文章，本文基于你给项目跑通，并 pip 集成过 flask 后

本修改支持加载一次模型，持续提供 HTTP 服务

在根目录新建主文件 run_serv.py

import io

from deoldify import device
from deoldify.device_id import DeviceId
from deoldify.visualize import *
from flask import Flask, request, Response

app = Flask(__name__)

# 提取通用代码
# choices:  CPU, GPU0...GPU7
device.set(device=DeviceId.GPU0)
plt.style.use('dark_background')
colorizer = get_image_colorizer(artistic=True)
torch.backends.cudnn.benchmark = True


@app.route('/action', methods=['POST', 'GET'])
def action():
    file = request.files['file']
    img_name = file.filename
    _, ext = os.path.splitext(img_name)
    img_in_memory = io.BytesIO(file.read())

	#此处的35应该是种子而不是参数 可以改成25或者什么的，影响似乎不大
    img_ret = colorizer.get_transformed_image_pzj(img_in_memory, 35)
    # colorizer.plot_transformed_image(path='C:/Users/PiaoZhenJia/Pictures/微信图片_20230728194252.png',
    #                                  render_factor=35, results_dir=Path('D:/'), post_process=False, watermarked=False)

    img_ret_bytes = io.BytesIO()
    img_ret.save(img_ret_bytes, format='PNG')
    img_ret_bytes.seek(0)
    res = Response(img_ret_bytes)
    res.headers.add('Content-Type', 'image/' + ext)
    res.headers.add('Content-Disposition', f'attachment; filename=colorize_{img_name}')
    return res

def main():
    app.run(port=2021, host="127.0.0.1", debug=False)


if __name__ == '__main__':
    main()

然后修改 visualize.py 我选择在187行附近加入这个函数

    def get_transformed_image_pzj(
            self, img_bytes: io.BytesIO, render_factor: int = None
    ) -> Image:
        self._clean_mem()
        orig_image = PIL.Image.open(img_bytes)
        filtered_image = self.filter.filter(
            orig_image, orig_image, render_factor=render_factor
        )

        return filtered_image