数据集准备

按照该链接下载coco数据集的指定类，作为训练数据集。YOLO-Coco-Dataset-Custom-Classes-Extractor

开源模型介绍

风格迁移模型参照该github开源项目。onnx_small_style 。

  • images/目录下包含输入图片文件夹、输出图片文件夹及风格图片文件夹
  • model/目录下有按特定风格图片及尺寸训练好的模型，可以直接使用
  • 风格迁移在线demo1：风格迁移网站
  • 风格迁移在线demo2：风格迁移网站

项目内容介绍

模型训练及测试

1
2
3
4
# 使用训练好的模型
python neural_style/neural_style.py eval --content-image images/content-images/birds.jpeg --model model/rain_princess.model --output-image images/output-images/birds_rain.jpg --cuda 0 # cpu
# 训练自己的模型
python neural_style/neural_style.py train --dataset dataset --style-image images/style-images/girl.jpg --save-model-dir model --epochs 2 --cuda 1 # gpu

注意事项

  1. 训练过程中报内存不足时，考虑风格图片的尺寸是否过大，或者减小batchsize
  2. 如果使用自己的笔记本训练数据，最好不要使用完整的coco数据集，可以只选取某些类别的图像数据