Skip to content

S-HuaBomb/pytorch-in-flask

Repository files navigation

风格转换

图像风格转换的深度学习算法研究和应用开发

[Live Demo: WCT Style Transfer] 部署在阿里云ECS,2020.10已经过期。很后悔忘了低价续费 😭

简介

实现一个高效的图像风格转换算法,通过Web应用让更多人使用图像风格转换技术进行娱乐或艺术创作。

算法

为了实现图像风格转换,参考了以下两篇论文,本应用的图像风格转换模块是基于第二篇论文的PyTorch实现:

项目结构

pytorch-in-flask/

	server:
	│  app_stylize.py  // flask服务端,调用wct的风格转换API
	│  Loader.py  // 图片加载方法
	│  modelsNIPS.py  // 编码器、解码器网络结构的定义
	│  requirements.txt  // 项目依赖清单
	│  torchfile.py  // 修改模型加载load()方法
	│  util.py  // 模型加载、WCT算法实现
	│  wct.py  // 风格转换API
	│
	├─models  // 用于风格转换的10个预训练模型
	│
	├─output
	│  ├─contents  // 存储用户上传的内容图
	│  └─stylized  // 存储风格化结果输出
	│
	├─styles  // 前端的九宫格风格图片原图
	
	src:
	│  App.vue  主vue组件
	│  main.js  // 前端项目的入口js文件
	│
	├─assets  // 前端用图
	│  ├─icons  // 存储.svg图片
	│  └─thumbs  // 存储前端九宫格中的风格图
	│
	├─components
	│      DrawingBoard.vue  // 画板组件,内容图展示在画板中
	│      ImageItem.vue  // 图片展示组件,用于遍历图片并展示
	│      LandingPage.vue  // 前端主页组件
	│
	└─router
	|        index.js  // 前端路由js文件
	|       
	index.html  // 入口html文件

本地安装

预训练模型

将这些训练好的Release中的模型放入项目中的 server/models 文件夹下。

CUDA

本人的笔记本开发环境如下:

项目 内容
CPU AMD 3550H, 2.10GHz
RAM 16GB
GPU AMD...
OS Windows10 (64bit)
如果可以配置CUDA 请配置如下:
CUDA CUDA 9.0 + cuDNN 7.4
PyTorch torch1.0.1

所以,在本机无法配置 CUDA 的硬伤下,本项目跑完一次风格转换的用时是60秒左右。 live demo 中的服务器配置页比较低,是最便宜的没有 GPU 的阿里云服务器。我去掉一层风格转换网络,并且将输入图片裁剪到 256×256 后才能达到 15 秒左右的速度。

CUDA能安装尽量安装。

后端

后端使用Python和Flask。只需写一个路由处理前后端通信即可。 进入到 server 目录,服务端代码和PyTorch风格转换代码都在里面,启动服务端的脚本文件是:

  • app_stylize.py

依赖

确保本机安装了Python3.7, 以下这些依赖是必须的: torch, pillow, flask, gevent。 可以直接通过 requirements.txt 来安装这些依赖:

pip install -r requirements.txt

运行

python app_stylize.py

运行成功的程序会在 5002 端口监听请求。浏览器进入这个URL:http://localhost:5002 就能看到 /index 首页路由的简单返回。

前端

需要安装好,配置好 npm 的系统环境:

# 克隆本项目到本地
git clone git@github.com:S-HuaBomb/pytorch-in-flask.git
cd pytorch-in-flask

# 安装前端依赖
npm install

# 运行前端
npm run dev

在浏览器打开这个URL:http://localhost:8080,就能看到前端界面,并且能跟后端完美通信。

部署上线

参考我的博客:Flask+Vue+PyTorch深度学习前后端分离部署,配合Gunicorn+Nginx+阿里云服务器