# text2image

**Repository Path**: jia_cheng_sui/text2image

## Basic Information

- **Project Name**: text2image
- **Description**: No description available
- **Primary Language**: Unknown
- **License**: Not specified
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2021-10-25
- **Last Updated**: 2021-12-01

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# text2image

作者：隋佳成、余乐

简介：文章配图工具 （本项目仅用于学习自然语言处理和torch，不用做其他用途）

## UpDate 2021.11.1
* 已有的数据集有coco、flower、bird，当然，如果你想试试自己的数据集，也没问题
* 当然了，我们不仅需要图片数据集，还需要其对应的embedding，这个embedding工作是我们自己做的，所以只需要下载图片数据集
* coco图片数据集对应的caption、还有我们自己得到的embedding数据都存放在了data文件夹下

| 图片数据集 | 对应地址 |
| :----- | ------  |
| coco | [地址](https://cocodataset.org/#download) |
| bird | [地址](http://www.vision.caltech.edu/visipedia/CUB-200-2011.html) |
| flower | [地址](https://www.robots.ox.ac.uk/~vgg/data/flowers/102/) |

## 项目描述
* 本项目是一个文章配图生成工具，基于StackGAN生成配图，输入一段文本，此工具会检测最优句子，并依此句子生成配图，默认生成两张配图。
* 本项目参考了StackGAN论文原文，复现其代码并做优化，最终将其适配成能够从文章中提取关键句子并生成配图的工具。
* 本项目通过Flask框架搭建了一个Web服务，将新闻摘要生成模型进行工程化，可以通过页面可视化地体验配图生成效果。
* 本项目代码中的关键之处全部有注释，可以自行查看。

## 文件结构
* config 保存模型各种配置文件的文件夹

* data 存放数据的文件夹

* templates 存放html页面的文件夹

* static 保存生成的配图和配图的位置、对应句子等信息的文件夹，网页从该文件夹读取信息

* main.py 顾名思义的就是最终的执行文件，其中可以读取命令行参数，并作初始化后进行训练或者加载模型后进行predict

   实际上我们用最新的main1.py文件

* model.py 文件中定义了各个组成StackGAN需要用到的网络结构，比如残差网络、阶段一生成器判别器等等一系列结构，每一个结构几乎都是继承了torch.nn.Module类。

* trainer.py 文件只包含一个类GANTrainer，用于指定训练的具体过程。

* utils.py文件中包含了一些需要用到的辅助方法，比如计算损失、KL度量等。当然了，这个文件中所包含的辅助方法是不够的，我们又还加了很多，新加的大多没有存在这个文件中。

* findMostPossiable.py 文件中的find方法执行了这个寻找最佳句子的功能，另外还有寻找最佳句子在文中位置，最佳句子对应图片的方法

* execute.py 一步化执行文件，用于执行测试与预测生成图片

* http_helper.py 用于网页可视化此配图生成工具，具体运行的时候，此文件会调用execute.py 并打开本地服务器，所以一般运行此文件

* save_sen2txt.py 文件中包含了一些函数，都用于展示效果时将输入的文本中选取的语句对应的图片、位置等等待展示的素材存放到static中，方便网页展示。

* my_pretreat.py 用于对数据进行初始化，从数据集文件中读取图片和caption，并将caption转化为embedding，在把图片和embedding按照一定规范以pickle形式存储，方便后续使用

* new_textdataset.py 文件主要定义了一个类TextDataset，用于存放和表达数据集的数据结构。区别于原论文的textdataset.py，由于采用的数据集不一致，对原数据结构稍作增删。

* pt_config.py 用于调整和修改一些超参数、训练、测试的参数等等。

## 运行环境
* pytorch == 1.9.1
* sentence_transformers == 2.1.0
* numpy == 1.19.5
* tensorboardX == 2.4
* easydict
* pandas

## 模型及训练参数
详细见miscc/pt_config.py文件，注释非常详细

## 模型文件分享
| 模型   | 存放地址 |
| ------  | -------------  |
| state1生成器 | output/Model/netG1 |
| state2生成器 | output/Model/netG2 |

## 模型训练
```
#阶段一训练 多GPU
python main1.py --cfg config/coco_s1.yml --gpu "0,1,3,4,5"
#阶段一训练 单GPU
python main1.py --cfg config/coco_s1.yml --gpu 0

#阶段二训练 多GPU
python main1.py --cfg config/coco_s2.yml --gpu "0,1,3,4,5"
#阶段二训练 单GPU
python main1.py --cfg config/coco_s2.yml --gpu 0
```
训练时间很长，不建议完整跑完。

## 模型测试

直接运行http_helper.py文件，然后浏览器打开 127.0.0.1:5555/graph-generate，从输入框中输入文本即可

## 未来工作
由于训练时间太长，我们利用coco数据集训练，5张2080显卡同时训练，训练一代阶段一需要20分钟左右，训练一代阶段二需要三小时左右，在现有实验室条件下不可能训练出很好结果来，所以，未来首先聚焦的工作肯定是将现有模型找机会跑完，看看效果，然后，再对网络结构进行改进看看效果。

- [ ] 利用coco数据集，完整运行阶段一阶段二训练1000代
- [ ] 利用bird数据集，完整运行阶段一阶段二训练120代
- [ ] 改进网络模型，比如改成deconv，再训练看效果
- [ ] 改进encoder，使embedding更加体现caption特征