DeOldify 接口化改造集成 Flask

本文主要是介绍DeOldify 接口化改造集成 Flask，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

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

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

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

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

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

import iofrom deoldify import device
from deoldify.device_id import DeviceId
from deoldify.visualize import *
from flask import Flask, request, Responseapp = 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 resdef 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