diff --git a/.gitignore b/.gitignore
index 973701941c8c07162d56b2656be17c5ea6bdf5e2..c1a7e93d334778204b805743b7469df3d456667a 100644
--- a/.gitignore
+++ b/.gitignore
@@ -138,7 +138,7 @@ mask_old/
origin_image/
datasets/
dataset/
-test*/
+test*
video_tmp/
result/
#./
@@ -163,14 +163,12 @@ result/
/make_datasets/datasets_img
/make_datasets/videos
#./models
-/models/videoHD_model.py
+#/models/videoHD_model.py
#./train
/train/clean/dataset
#mediafile
*iter
*.pth
-*.png
-*.jpg
*.jpeg
*.bmp
*.mp4
diff --git a/README.md b/README.md
index 15977a2819e6f1e469117d76a40e5aed069ef58f..0aa37c9b8719993da38795058f571d23175a3838 100755
--- a/README.md
+++ b/README.md
@@ -1,44 +1,49 @@
# 
DeepMosaics
You can use it to automatically remove the mosaics in images and videos, or add mosaics to them.
-This porject based on ‘semantic segmentation’ and ‘Image-to-Image Translation’.
-Master is not stable. Please use a [stable version](https://github.com/HypoX64/DeepMosaics/tree/stable)
-* [中文版](./README_CN.md)
+This porject based on "semantic segmentation" and "Image-to-Image Translation".
+
+* [中文版README](./README_CN.md)
### More example
origin | auto add mosaic | auto clean mosaic
:-:|:-:|:-:
- |  | 
- |  | 
+ |  | 
+ |  | 
* Compared with [DeepCreamPy](https://github.com/deeppomf/DeepCreamPy)
-
mosaic image | DeepCreamPy | ours
:-:|:-:|:-:
- |  | 
- |  | 
+ |  | 
+ |  | 
+* Style Transfer
+origin | to Van Gogh | to winter
+:-:|:-:|:-:
+ |  | 
+An interesting example:[Ricardo Milos to cat](https://www.bilibili.com/video/BV1Q7411W7n6)
## Run DeepMosaics
You can either run DeepMosaics via pre-built binary package or from source.
### Pre-built binary package
For windows, we bulid a GUI version for easy test.
-Download this version via [[Google Drive]](https://drive.google.com/open?id=1LTERcN33McoiztYEwBxMuRjjgxh4DEPs) [[百度云,提取码1x0a]](https://pan.baidu.com/s/10rN3U3zd5TmfGpO_PEShqQ)
+Download this version and pre-trained model via [[Google Drive]](https://drive.google.com/open?id=1LTERcN33McoiztYEwBxMuRjjgxh4DEPs) [[百度云,提取码1x0a]](https://pan.baidu.com/s/10rN3U3zd5TmfGpO_PEShqQ)
+* [[How to use]](./docs/exe_help.md)
-
Attentions:
+
- Require Windows_x86_64, Windows10 is better.
- - Different pre-trained models are suitable for different effects.
- - Run time depends on computer performance.
- - If output video cannot be played, you can try with [potplayer](https://daumpotplayer.com/download/).
- - GUI version update slower than source.
+ - Different pre-trained models are suitable for different effects.[[Introduction to pre-trained models]](./docs/pre-trained_models_introduction.md)
+ - Run time depends on computer performance(The current version does not support gpu, if you need to use gpu please run source).
+ - If output video cannot be played, you can try with [potplayer](https://daumpotplayer.com/download/).
+ - GUI version update slower than source.
### Run from source
#### Prerequisites
- Linux, Mac OS, Windows
- Python 3.6+
- [ffmpeg 3.4.6](http://ffmpeg.org/)
- - [Pytorch 1.0+](https://pytorch.org/) [(Old version codes)](https://github.com/HypoX64/DeepMosaics/tree/Pytorch0.4)
+ - [Pytorch 1.0+](https://pytorch.org/)
- CPU or NVIDIA GPU + CUDA CuDNN
#### Dependencies
This code depends on opencv-python, torchvision available via pip install.
@@ -47,22 +52,23 @@ This code depends on opencv-python, torchvision available via pip install.
git clone https://github.com/HypoX64/DeepMosaics
cd DeepMosaics
```
-#### Get pre_trained models and test video
+#### Get pre-trained models
You can download pre_trained models and put them into './pretrained_models'.
-[[Google Drive]](https://drive.google.com/open?id=1LTERcN33McoiztYEwBxMuRjjgxh4DEPs) [[百度云,提取码1x0a]](https://pan.baidu.com/s/10rN3U3zd5TmfGpO_PEShqQ)
+[[Google Drive]](https://drive.google.com/open?id=1LTERcN33McoiztYEwBxMuRjjgxh4DEPs) [[百度云,提取码1x0a]](https://pan.baidu.com/s/10rN3U3zd5TmfGpO_PEShqQ)
+[[Introduction to pre-trained models]](./docs/pre-trained_models_introduction.md)
#### Simple example
-* Add Mosaic (output video will save in './result')
+* Add Mosaic (output video will save in './result')
```bash
python3 deepmosaic.py --media_path ./imgs/ruoruo.jpg --model_path ./pretrained_models/mosaic/add_face.pth --use_gpu -1
```
-* Clean Mosaic (output video will save in './result')
+* Clean Mosaic (output video will save in './result')
```bash
python3 deepmosaic.py --media_path ./result/ruoruo_add.jpg --model_path ./pretrained_models/mosaic/clean_face_HD.pth --use_gpu -1
```
#### More parameters
-If you want to test other image or video, please refer to this file.
-[[options.py]](./cores/options.py)
+If you want to test other image or video, please refer to this file.
+[[options_introduction.md]](./docs/options_introduction.md)
## Acknowledgments
This code borrows heavily from [[pytorch-CycleGAN-and-pix2pix]](https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix) [[Pytorch-UNet]](https://github.com/milesial/Pytorch-UNet)[[pix2pixHD]](https://github.com/NVIDIA/pix2pixHD).
diff --git a/README_CN.md b/README_CN.md
index b4bef9bc188dad0d4d5eb3251e944c36b2707cb4..d9f9c4c45005be129181ad7444ad34103c4b27ed 100644
--- a/README_CN.md
+++ b/README_CN.md
@@ -1,23 +1,22 @@
# DeepMosaics
这是一个通过深度学习自动的为图片/视频添加马赛克,或消除马赛克的项目.
它基于“语义分割”以及“图像翻译”.
-主分支并不稳定,请移步[稳定版本](https://github.com/HypoX64/DeepMosaics/tree/stable)
+
### 更多例子
原始 | 自动打码 | 自动去码
:-:|:-:|:-:
- |  | 
- |  | 
+ |  | 
+ |  | 
* 与 [DeepCreamPy](https://github.com/deeppomf/DeepCreamPy)相比较
-
马赛克图片 | DeepCreamPy | ours
:-:|:-:|:-:
- |  | 
- |  | 
-
-## 一些说明
-* 训练部分并不完全
-* 现在,代码已经支持基于[pix2pixHD](https://github.com/NVIDIA/pix2pixHD)训练出的模型,但网络仍在训练中,这将使得输出结果看起来更加清晰,"真实".
-* 新的模型,可根据视频帧间关系进行马赛克恢复,在pretrained model 中被命名为*_video.pth
+ |  | 
+ |  | 
+* 风格转换
+原始 | 梵高风格 | 转化为冬天
+:-:|:-:|:-:
+ |  | 
+一个有意思的尝试:[香蕉君♂猫](https://www.bilibili.com/video/BV1Q7411W7n6)
## 如何运行
可以通过我们预编译好的二进制包或源代码运行.
@@ -26,12 +25,14 @@
对于Windows用户,我们提供了包含GUI界面的免安装软件包.
可以通过下面两种方式进行下载: [[Google Drive]](https://drive.google.com/open?id=1LTERcN33McoiztYEwBxMuRjjgxh4DEPs) [[百度云,提取码1x0a]](https://pan.baidu.com/s/10rN3U3zd5TmfGpO_PEShqQ)
+* [[使用教程]](./docs/exe_help_CN.md)
+
注意事项:
- 程序的运行要求在64位Windows操作系统,我仅在Windows10运行过,其他版本暂未经过测试
- - 请根据需求选择合适的预训练模型进行测试
- - 运行时间取决于电脑性能,对于视频文件,我们建议可以先使用截图进行测试.
+ - 请根据需求选择合适的预训练模型进行测试，不同的预期训练模型具有不同的效果.[[预训练模型介绍]](./docs/pre-trained_models_introduction_CN.md)
+ - 运行时间取决于电脑性能,对于视频文件,我们建议使用源码并在GPU上运行.
- 如果输出的视频无法播放,这边建议您尝试[potplayer](https://daumpotplayer.com/download/).
- 相比于源码,该版本的更新将会延后.
@@ -40,7 +41,7 @@
- Linux, Mac OS, Windows
- Python 3.6+
- [ffmpeg 3.4.6](http://ffmpeg.org/)
- - [Pytorch 1.0+](https://pytorch.org/) [(Old version codes)](https://github.com/HypoX64/DeepMosaics/tree/Pytorch0.4)
+ - [Pytorch 1.0+](https://pytorch.org/)
- CPU or NVIDIA GPU + CUDA CuDNN
#### Python依赖项
代码依赖于opencv-python以及 torchvision,可有通过pip install 进行安装.
@@ -52,19 +53,20 @@ cd DeepMosaics
#### 下载预训练模型
可以通过以下两种方法下载预训练模型,并将他们置于'./pretrained_models'文件夹中.
[[Google Drive]](https://drive.google.com/open?id=1LTERcN33McoiztYEwBxMuRjjgxh4DEPs) [[百度云,提取码1x0a]](https://pan.baidu.com/s/10rN3U3zd5TmfGpO_PEShqQ)
+[[预训练模型介绍]](./docs/pre-trained_models_introduction_CN.md)
#### 简单的例子
-* 为视频添加马赛克,例子中认为脸是需要打码的区域 ,可以通过切换预训练模型切换自动打码区域(输出结果将储存到 './result')
+* 为视频添加马赛克,例子中认为脸是需要打码的区域 ,可以通过切换预训练模型切换自动打码区域(输出结果将储存到 './result')
```bash
python3 deepmosaic.py --media_path ./imgs/ruoruo.jpg --model_path ./pretrained_models/mosaic/add_face.pth --use_gpu -1
```
-* 将视频中的马赛克移除,对于不同的打码物体需要使用对应的预训练模型进行马赛克消除(输出结果将储存到 './result')
+* 将视频中的马赛克移除,对于不同的打码物体需要使用对应的预训练模型进行马赛克消除(输出结果将储存到 './result')
```bash
python3 deepmosaic.py --media_path ./result/ruoruo_add.jpg --model_path ./pretrained_models/mosaic/clean_face_HD.pth --use_gpu -1
```
#### 更多的参数
-如果想要测试其他的图片或视频,请参照以下文件输入参数.
-[[options.py]](./cores/options.py)
+如果想要测试其他的图片或视频,请参照以下文件输入参数.
+[[options_introduction_CN.md]](./docs/options_introduction_CN.md)
## 鸣谢
-代码大量的参考了以下项目:[[pytorch-CycleGAN-and-pix2pix]](https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix) [[Pytorch-UNet]](https://github.com/milesial/Pytorch-UNet)[[pix2pixHD]](https://github.com/NVIDIA/pix2pixHD)..
+代码大量的参考了以下项目:[[pytorch-CycleGAN-and-pix2pix]](https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix) [[Pytorch-UNet]](https://github.com/milesial/Pytorch-UNet)[[pix2pixHD]](https://github.com/NVIDIA/pix2pixHD).
\ No newline at end of file
diff --git a/cores/core.py b/cores/core.py
index e5ed377cf784d0f0d8c437eaf2d673d6c45df209..22a408a0cc6dec1663f83d4eff8ebe2e242356a8 100644
--- a/cores/core.py
+++ b/cores/core.py
@@ -150,7 +150,10 @@ def cleanmosaic_video_byframe(opt,netG,netM):
def cleanmosaic_video_fusion(opt,netG,netM):
path = opt.media_path
N = 25
- INPUT_SIZE = 128
+ if 'HD' in os.path.basename(opt.model_path):
+ INPUT_SIZE = 256
+ else:
+ INPUT_SIZE = 128
fps,imagepaths,height,width = video_init(opt,path)
positions = get_mosaic_positions(opt,netM,imagepaths,savemask=True)
diff --git a/cores/options.py b/cores/options.py
index cf31aacf7fb1a2319e89a29f6be75635a2a8865d..b3cb051608ad3fcf7e2e2786e2647e7cef8f01d3 100644
--- a/cores/options.py
+++ b/cores/options.py
@@ -13,36 +13,36 @@ class Options():
self.parser.add_argument('--use_gpu',type=int,default=0, help='if -1, do not use gpu')
# self.parser.add_argument('--use_gpu', action='store_true', help='if input it, use gpu')
self.parser.add_argument('--media_path', type=str, default='./imgs/ruoruo.jpg',help='your videos or images path')
- self.parser.add_argument('--mode', type=str, default='auto',help='auto | add | clean | style')
+ self.parser.add_argument('--mode', type=str, default='auto',help='Program running mode. auto | add | clean | style')
self.parser.add_argument('--model_path', type=str, default='./pretrained_models/mosaic/add_face.pth',help='pretrained model path')
self.parser.add_argument('--result_dir', type=str, default='./result',help='output media will be saved here')
self.parser.add_argument('--tempimage_type', type=str, default='png',help='type of temp image, png | jpg, png is better but occupy more storage space')
self.parser.add_argument('--netG', type=str, default='auto',
help='select model to use for netG(Clean mosaic and Transfer style) -> auto | unet_128 | unet_256 | resnet_9blocks | HD | video')
self.parser.add_argument('--fps', type=int, default=0,help='read and output fps, if 0-> origin')
- self.parser.add_argument('--output_size', type=int, default=0,help='size of output file,if 0 -> origin')
+ self.parser.add_argument('--output_size', type=int, default=0,help='size of output media, if 0 -> origin')
#AddMosaic
self.parser.add_argument('--mosaic_mod', type=str, default='squa_avg',help='type of mosaic -> squa_avg | squa_random | squa_avg_circle_edge | rect_avg | random')
self.parser.add_argument('--mosaic_size', type=int, default=0,help='mosaic size,if 0 auto size')
- self.parser.add_argument('--mask_extend', type=int, default=10,help='more mosaic area')
+ self.parser.add_argument('--mask_extend', type=int, default=10,help='extend mosaic area')
self.parser.add_argument('--mask_threshold', type=int, default=64,help='threshold of recognize mosaic position 0~255')
#CleanMosaic
self.parser.add_argument('--mosaic_position_model_path', type=str, default='auto',help='name of model use to find mosaic position')
- self.parser.add_argument('--traditional', action='store_true', help='if true, use traditional image processing methods to clean mosaic')
+ self.parser.add_argument('--traditional', action='store_true', help='if specified, use traditional image processing methods to clean mosaic')
self.parser.add_argument('--tr_blur', type=int, default=10, help='ksize of blur when using traditional method, it will affect final quality')
self.parser.add_argument('--tr_down', type=int, default=10, help='downsample when using traditional method,it will affect final quality')
- self.parser.add_argument('--no_feather', action='store_true', help='if true, no edge feather and color correction, but run faster')
- self.parser.add_argument('--all_mosaic_area', action='store_true', help='if true, find all mosaic area, else only find the largest area')
+ self.parser.add_argument('--no_feather', action='store_true', help='if specified, no edge feather and color correction, but run faster')
+ self.parser.add_argument('--all_mosaic_area', action='store_true', help='if specified, find all mosaic area, else only find the largest area')
self.parser.add_argument('--medfilt_num', type=int, default=11,help='medfilt window of mosaic movement in the video')
self.parser.add_argument('--ex_mult', type=str, default='auto',help='mosaic area expansion')
#StyleTransfer
self.parser.add_argument('--preprocess', type=str, default='resize', help='resize and cropping of images at load time [ resize | resize_scale_width | edges | gray] or resize,edges(use comma to split)')
- self.parser.add_argument('--edges', action='store_true', help='if true, use edges to generate pictures,(input_nc = 1)')
+ self.parser.add_argument('--edges', action='store_true', help='if specified, use edges to generate pictures,(input_nc = 1)')
self.parser.add_argument('--canny', type=int, default=150,help='threshold of canny')
- self.parser.add_argument('--only_edges', action='store_true', help='if true, output media will be edges')
+ self.parser.add_argument('--only_edges', action='store_true', help='if specified, output media will be edges')
self.initialized = True
@@ -67,7 +67,9 @@ class Options():
elif 'style' in model_name or 'edges' in model_name:
self.opt.mode = 'style'
else:
- print('Please input running mode!')
+ print('Please input running model!')
+ input('Please press any key to exit.\n')
+ exit(0)
if self.opt.output_size == 0 and self.opt.mode == 'style':
self.opt.output_size = 512
@@ -80,12 +82,14 @@ class Options():
self.opt.netG = 'unet_128'
elif 'resnet_9blocks' in model_name:
self.opt.netG = 'resnet_9blocks'
- elif 'HD' in model_name:
+ elif 'HD' in model_name and 'video' not in model_name:
self.opt.netG = 'HD'
elif 'video' in model_name:
self.opt.netG = 'video'
else:
print('Type of Generator error!')
+ input('Please press any key to exit.\n')
+ exit(0)
if self.opt.ex_mult == 'auto':
if 'face' in model_name:
diff --git a/deepmosaic.py b/deepmosaic.py
index dec510296deeb08d011adcaa4ea064a9d9782039..58a6031b86742ce458ecab61012a2c42a58fbef9 100644
--- a/deepmosaic.py
+++ b/deepmosaic.py
@@ -1,4 +1,6 @@
import os
+import sys
+import traceback
from cores import Options,core
from util import util
from models import loadmodel
@@ -58,12 +60,18 @@ def main():
util.clean_tempfiles(tmp_init = False)
-main()
-# if __name__ == '__main__':
-# try:
-# main()
-# except Exception as e:
-# print('Error:',e)
-# input('Please press any key to exit.\n')
-# util.clean_tempfiles(tmp_init = False)
-# exit(0)
+# main()
+if __name__ == '__main__':
+ try:
+ main()
+ except Exception as ex:
+ print('--------------------ERROR--------------------')
+ ex_type, ex_val, ex_stack = sys.exc_info()
+ print('Error Type:',ex_type)
+ print(ex_val)
+ for stack in traceback.extract_tb(ex_stack):
+ print(stack)
+ input('Please press any key to exit.\n')
+ util.clean_tempfiles(tmp_init = False)
+ exit(0)
+
diff --git a/docs/exe_help.md b/docs/exe_help.md
new file mode 100644
index 0000000000000000000000000000000000000000..1b5bc09f8cbcd57b76b99d1ba8fd1be9d6a7d97a
--- /dev/null
+++ b/docs/exe_help.md
@@ -0,0 +1,86 @@
+## DeepMosaics.exe Instructions
+[[中文版]](./exe_help_CN.md)
+This is a GUI version compiled in Windows.
+Download this version and pre-trained model via [[Google Drive]](https://drive.google.com/open?id=1LTERcN33McoiztYEwBxMuRjjgxh4DEPs) [[百度云,提取码1x0a]](https://pan.baidu.com/s/10rN3U3zd5TmfGpO_PEShqQ)
+Attentions:
+
+ - Require Windows_x86_64, Windows10 is better.
+ - Different pre-trained models are suitable for different effects.
+ - Run time depends on computer performance.
+ - If output video cannot be played, you can try with [potplayer](https://daumpotplayer.com/download/).
+ - GUI version update slower than source.
+### How to use
+* step 1: Choose image or video.
+* step 2: Choose model(Different pre-trained models are suitable for different effects)
+* step3: Run program and wait.
+* step4: Cheek reult in './result'.
+
+### Introduction to pre-trained models
+* Mosaic
+| Name | Description |
+| :------------------------------: | :---------------------------------------------------------: |
+| add_face.pth | Add mosaic to all faces in images/videos. |
+| clean_face_HD.pth | Clean mosaic to all faces in images/video.
(RAM > 8GB). |
+| add_youknow.pth | Add mosaic to all (FBI Warning) in images/videos. |
+| clean_youknow_resnet_9blocks.pth | Clean mosaic to all (FBI Warning) in images/videos. |
+| clean_youknow_video.pth | Clean mosaic to all (FBI Warning) in videos. |
+| clean_youknow_video_HD.pth | Clean mosaic to all (FBI Warning) in videos.
(RAM > 8GB) |
+
+* Style Transfer
+| Name | Description |
+| :---------------------: | :-------------------------------------------------------: |
+| style_apple2orange.pth | Convert apples to oranges. |
+| style_orange2apple.pth | Convert oranges to apples |
+| style_summer2winter.pth | Convert summer to winter. |
+| style_winter2summer.pth | Convert winter to summer. |
+| style_cezanne.pth | Convert photos/video to Paul Cézanne style. |
+| style_monet.pth | Convert photos/video to Claude Monet style. |
+| style_ukiyoe.pth | Convert photos/video to Ukiyoe style. |
+| style_vangogh.pth | Convert photos/video to Van Gogh style. |
+### Annotation
+
+* 1. Choose image or video.
+* 2. Choose model(Different pre-trained models are suitable for different effects).
+* 3. Program running mode. (auto | add | clean | style)
+* 4. Use GPU to run deep learning model. (The current version does not support gpu, if you need to use gpu please run source).
+* 5. Limit the fps of the output video(0->original fps).
+* 6. More options.
+* 7. More options can be input.
+* 8. Run program.
+* 9. Open help file.
+* 10. Sponsor our project.
+* 11. Version information.
+* 12. Open the URL on github.
+
+### Introduction to options
+If you need more effects, use '--option your-parameters' to enter what you need.
+* Base
+| Option | Description | Default |
+| :----------: | :----------------------------------------: | :-------------------------------------: |
+| --use_gpu | if -1, do not use gpu | 0 |
+| --media_path | your videos or images path | ./imgs/ruoruo.jpg |
+| --mode | program running mode(auto/clean/add/style) | 'auto' |
+| --model_path | pretrained model path | ./pretrained_models/mosaic/add_face.pth |
+| --result_dir | output media will be saved here | ./result |
+| --fps | read and output fps, if 0-> origin | 0 |
+
+* AddMosaic
+| Option | Description | Default |
+| :--------------: | :----------------------------------------------------------: | :------: |
+| --mosaic_mod | type of mosaic -> squa_avg/ squa_random/ squa_avg_circle_edge/ rect_avg/random | squa_avg |
+| --mosaic_size | mosaic size,if 0 -> auto size | 0 |
+| --mask_extend | extend mosaic area | 10 |
+| --mask_threshold | threshold of recognize mosaic position 0~255 | 64 |
+
+* CleanMosaic
+| Option | Description | Default |
+| :-----------: | :----------------------------------------------------------: | :-----: |
+| --traditional | if specified, use traditional image processing methods to clean mosaic | |
+| --tr_blur | ksize of blur when using traditional method, it will affect final quality | 10 |
+| --tr_down | downsample when using traditional method,it will affect final quality | 10 |
+| --medfilt_num | medfilt window of mosaic movement in the video | 11 |
+
+* Style Transfer
+| Option | Description | Default |
+| :-----------: | :----------------------------------: | :-----: |
+| --output_size | size of output media, if 0 -> origin | 512 |
\ No newline at end of file
diff --git a/docs/exe_help_CN.md b/docs/exe_help_CN.md
new file mode 100644
index 0000000000000000000000000000000000000000..69624d48f2f1a5386f5567e49d0845386696656a
--- /dev/null
+++ b/docs/exe_help_CN.md
@@ -0,0 +1,85 @@
+## DeepMosaics.exe 使用说明
+下载程序以及预训练模型 [[Google Drive]](https://drive.google.com/open?id=1LTERcN33McoiztYEwBxMuRjjgxh4DEPs) [[百度云,提取码1x0a]](https://pan.baidu.com/s/10rN3U3zd5TmfGpO_PEShqQ)
+注意事项:
+
+ - 程序的运行要求在64位Windows操作系统,我仅在Windows10运行过,其他版本暂未经过测试
+ - 请根据需求选择合适的预训练模型进行测试
+ - 运行时间取决于电脑性能,对于视频文件,我们建议使用源码以及GPU运行
+ - 如果输出的视频无法播放,这边建议您尝试[potplayer](https://daumpotplayer.com/download/).
+ - 相比于源码,该版本的更新将会延后.
+
+### 如何使用
+* step 1: 选择需要处理的图片或视频
+* step 2: 选择预训练模型(不同的预训练模型有不同的效果)
+* step3: 运行程序并等待
+* step4: 查看结果(储存在result文件夹下)
+
+## 预训练模型说明
+当前的预训练模型分为两类——添加/移除马赛克以及风格转换.
+
+* 马赛克
+| 文件名 | 描述 |
+| :------------------------------: | :-------------------------------------------: |
+| add_face.pth | 对图片或视频中的脸部打码 |
+| clean_face_HD.pth | 对图片或视频中的脸部去码
(要求内存 > 8GB). |
+| add_youknow.pth | 对图片或视频中的十八禁内容打码 |
+| clean_youknow_resnet_9blocks.pth | 对图片或视频中的十八禁内容去码 |
+| clean_youknow_video.pth | 对视频中的十八禁内容去码 |
+| clean_youknow_video_HD.pth | 对视频中的十八禁内容去码
(要求内存 > 8GB) |
+
+* 风格转换
+| 文件名 | 描述 |
+| :---------------------: | :-------------------------------------------------------: |
+| style_apple2orange.pth | 苹果变橙子 |
+| style_orange2apple.pth | 橙子变苹果 |
+| style_summer2winter.pth | 夏天变冬天 |
+| style_winter2summer.pth | 冬天变夏天 |
+| style_cezanne.pth | 转化为Paul Cézanne 的绘画风格 |
+| style_monet.pth | 转化为Claude Monet的绘画风格 |
+| style_ukiyoe.pth | 转化为Ukiyoe的绘画风格 |
+| style_vangogh.pth | 转化为Van Gogh的绘画风格 |
+
+### GUI界面注释
+
+* 1. 选择需要处理的图片或视频
+* 2. 选择预训练模型
+* 3. 程序运行模式 (auto | add | clean | style)
+* 4. 使用GPU (该版本目前不支持GPU,若需要使用GPU请使用源码运行).
+* 5. 限制输出的视频帧率(0->原始帧率).
+* 6. 更多的选项以及参数
+* 7. 自行输入更多参数，详见下文
+* 8. 运行
+* 9. 打开帮助文件
+* 10. 支持我们
+* 11. 版本信息
+* 12. 打开项目的github页面
+
+### 参数说明
+如果需要更多的效果, 请按照 '--option your-parameters' 输入所需要的参数
+* 基本
+| 选项 | 描述 | 默认 |
+| :----------: | :------------------------: | :-------------------------------------: |
+| --use_gpu | if -1, do not use gpu | 0 |
+| --media_path | 需要处理的视频或者照片的路径 | ./imgs/ruoruo.jpg |
+| --mode | 运行模式(auto/clean/add/style) | 'auto' |
+| --model_path | 预训练模型的路径 | ./pretrained_models/mosaic/add_face.pth |
+| --result_dir | 保存路径 | ./result |
+| --fps | 限制视频输出的fps，0则为默认 | 0 |
+* 添加马赛克
+| 选项 | 描述 | 默认 |
+| :----------: | :------------------------: | :-------------------------------------: |
+| --mosaic_mod | 马赛克类型 -> squa_avg/ squa_random/ squa_avg_circle_edge/ rect_avg/random | squa_avg |
+| --mosaic_size | 马赛克大小，０则为自动 | 0 |
+| --mask_extend | 拓展马赛克区域 | 10 |
+| --mask_threshold | 马赛克区域识别阈值 0~255 | 64 |
+* 去除马赛克
+| 选项 | 描述 | 默认 |
+| :----------: | :------------------------: | :-------------------------------------: |
+| --traditional | 如果输入这个参数则使用传统方法清除马赛克 | |
+| --tr_blur | 传统方法模糊尺寸 | 10 |
+| --tr_down | 传统方法下采样尺寸 | 10 |
+| --medfilt_num | medfilt window of mosaic movement in the video | 11 |
+* 风格转换
+| 选项 | 描述 | 默认 |
+| :----------: | :------------------------: | :-------------------------------------: |
+| --output_size | 输出媒体的尺寸，如果是０则为原始尺寸 |512|
\ No newline at end of file
diff --git a/docs/options_introduction.md b/docs/options_introduction.md
new file mode 100644
index 0000000000000000000000000000000000000000..c8cfa66885445c9f12ce0fd204d7c25b58264304
--- /dev/null
+++ b/docs/options_introduction.md
@@ -0,0 +1,33 @@
+## Introduction to options
+If you need more effects, use '--option your-parameters' to enter what you need.
+
+### Base
+| Option | Description | Default |
+| :----------: | :------------------------: | :-------------------------------------: |
+| --use_gpu | if -1, do not use gpu | 0 |
+| --media_path | your videos or images path | ./imgs/ruoruo.jpg |
+| --mode | program running mode(auto/clean/add/style) | 'auto' |
+| --model_path | pretrained model path | ./pretrained_models/mosaic/add_face.pth |
+| --result_dir | output media will be saved here| ./result |
+| --fps | read and output fps, if 0-> origin | 0 |
+
+### AddMosaic
+| Option | Description | Default |
+| :----------: | :------------------------: | :-------------------------------------: |
+| --mosaic_mod | type of mosaic -> squa_avg/ squa_random/ squa_avg_circle_edge/ rect_avg/random | squa_avg |
+| --mosaic_size | mosaic size,if 0 -> auto size | 0 |
+| --mask_extend | extend mosaic area | 10 |
+| --mask_threshold | threshold of recognize mosaic position 0~255 | 64 |
+
+### CleanMosaic
+| Option | Description | Default |
+| :----------: | :------------------------: | :-------------------------------------: |
+| --traditional | if specified, use traditional image processing methods to clean mosaic | |
+| --tr_blur | ksize of blur when using traditional method, it will affect final quality | 10 |
+| --tr_down | downsample when using traditional method,it will affect final quality | 10 |
+| --medfilt_num | medfilt window of mosaic movement in the video | 11 |
+
+### Style Transfer
+| Option | Description | Default |
+| :----------: | :------------------------: | :-------------------------------------: |
+| --output_size | size of output media, if 0 -> origin |512|
\ No newline at end of file
diff --git a/docs/options_introduction_CN.md b/docs/options_introduction_CN.md
new file mode 100644
index 0000000000000000000000000000000000000000..f40f4884ea7ac79419215e47a60c5a72690eeb76
--- /dev/null
+++ b/docs/options_introduction_CN.md
@@ -0,0 +1,33 @@
+## 参数说明
+如果需要更多的效果, 请按照 '--option your-parameters' 输入所需要的参数
+
+### 基本
+| 选项 | 描述 | 默认 |
+| :----------: | :------------------------: | :-------------------------------------: |
+| --use_gpu | if -1, do not use gpu | 0 |
+| --media_path | 需要处理的视频或者照片的路径 | ./imgs/ruoruo.jpg |
+| --mode | 运行模式(auto/clean/add/style) | 'auto' |
+| --model_path | 预训练模型的路径 | ./pretrained_models/mosaic/add_face.pth |
+| --result_dir | 保存路径 | ./result |
+| --fps | 限制视频输出的fps，0则为默认 | 0 |
+
+### 添加马赛克
+| 选项 | 描述 | 默认 |
+| :----------: | :------------------------: | :-------------------------------------: |
+| --mosaic_mod | 马赛克类型 -> squa_avg/ squa_random/ squa_avg_circle_edge/ rect_avg/random | squa_avg |
+| --mosaic_size | 马赛克大小，０则为自动 | 0 |
+| --mask_extend | 拓展马赛克区域 | 10 |
+| --mask_threshold | 马赛克区域识别阈值 0~255 | 64 |
+
+### 去除马赛克
+| 选项 | 描述 | 默认 |
+| :----------: | :------------------------: | :-------------------------------------: |
+| --traditional | 如果输入这个参数则使用传统方法清除马赛克 | |
+| --tr_blur | 传统方法模糊尺寸 | 10 |
+| --tr_down | 传统方法下采样尺寸 | 10 |
+| --medfilt_num | medfilt window of mosaic movement in the video | 11 |
+
+### 风格转换
+| 选项 | 描述 | 默认 |
+| :----------: | :------------------------: | :-------------------------------------: |
+| --output_size | 输出媒体的尺寸，如果是０则为原始尺寸 |512|
\ No newline at end of file
diff --git a/docs/pre-trained_models_introduction.md b/docs/pre-trained_models_introduction.md
new file mode 100644
index 0000000000000000000000000000000000000000..a74355dab7f2343eed5402c3d8dbac399d056899
--- /dev/null
+++ b/docs/pre-trained_models_introduction.md
@@ -0,0 +1,26 @@
+## Introduction to pre-trained models
+The current pre-trained models are divided into two categories(Add/Clean mosaic and StyleTransfer).
+Download pre-trained model via [[Google Drive]](https://drive.google.com/open?id=1LTERcN33McoiztYEwBxMuRjjgxh4DEPs) [[百度云,提取码1x0a]](https://pan.baidu.com/s/10rN3U3zd5TmfGpO_PEShqQ)
+
+### Mosaic
+| Name | Description |
+| :------------------------------: | :---------------------------------------------------------: |
+| add_face.pth | Add mosaic to all faces in images/videos. |
+| clean_face_HD.pth | Clean mosaic to all faces in images/video.
(RAM > 8GB). |
+| add_youknow.pth | Add mosaic to all (FBI Warning) in images/videos. |
+| clean_youknow_resnet_9blocks.pth | Clean mosaic to all (FBI Warning) in images/videos. |
+| clean_youknow_video.pth | Clean mosaic to all (FBI Warning) in videos. |
+| clean_youknow_video_HD.pth | Clean mosaic to all (FBI Warning) in videos.
(RAM > 8GB) |
+
+### Style Transfer
+| Name | Description |
+| :---------------------: | :-------------------------------------------------------: |
+| style_apple2orange.pth | Convert apples to oranges. |
+| style_orange2apple.pth | Convert oranges to apples |
+| style_summer2winter.pth | Convert summer to winter. |
+| style_winter2summer.pth | Convert winter to summer. |
+| style_cezanne.pth | Convert photos/video to Paul Cézanne style. |
+| style_monet.pth | Convert photos/video to Claude Monet style. |
+| style_ukiyoe.pth | Convert photos/video to Ukiyoe style. |
+| style_vangogh.pth | Convert photos/video to Van Gogh style. |
+
diff --git a/docs/pre-trained_models_introduction_CN.md b/docs/pre-trained_models_introduction_CN.md
new file mode 100644
index 0000000000000000000000000000000000000000..95162d827fb5d8890f0a82670f5589bdc3b546c7
--- /dev/null
+++ b/docs/pre-trained_models_introduction_CN.md
@@ -0,0 +1,26 @@
+## 预训练模型说明
+当前的预训练模型分为两类——添加/移除马赛克以及风格转换.
+可以通过以下方式下载预训练模型 [[Google Drive]](https://drive.google.com/open?id=1LTERcN33McoiztYEwBxMuRjjgxh4DEPs) [[百度云,提取码1x0a]](https://pan.baidu.com/s/10rN3U3zd5TmfGpO_PEShqQ)
+
+### 马赛克
+| 文件名 | 描述 |
+| :------------------------------: | :-------------------------------------------: |
+| add_face.pth | 对图片或视频中的脸部打码 |
+| clean_face_HD.pth | 对图片或视频中的脸部去码
(要求内存 > 8GB). |
+| add_youknow.pth | 对图片或视频中的十八禁内容打码 |
+| clean_youknow_resnet_9blocks.pth | 对图片或视频中的十八禁内容去码 |
+| clean_youknow_video.pth | 对视频中的十八禁内容去码 |
+| clean_youknow_video_HD.pth | 对视频中的十八禁内容去码
(要求内存 > 8GB) |
+
+### 风格转换
+| 文件名 | 描述 |
+| :---------------------: | :-------------------------------------------------------: |
+| style_apple2orange.pth | 苹果变橙子 |
+| style_orange2apple.pth | 橙子变苹果 |
+| style_summer2winter.pth | 夏天变冬天 |
+| style_winter2summer.pth | 冬天变夏天 |
+| style_cezanne.pth | 转化为Paul Cézanne 的绘画风格 |
+| style_monet.pth | 转化为Claude Monet的绘画风格 |
+| style_ukiyoe.pth | 转化为Ukiyoe的绘画风格 |
+| style_vangogh.pth | 转化为Van Gogh的绘画风格 |
+
diff --git a/imgs/GUI.png b/imgs/GUI.png
index e9ffe2423665e63eb4db0e2e1e040c0dd3883508..7404e3bb52995fe04c352965f75f0ff35535c69e 100644
Binary files a/imgs/GUI.png and b/imgs/GUI.png differ
diff --git a/imgs/GUI_Instructions.jpg b/imgs/GUI_Instructions.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..600815241b6a3d3e4bfe5465af0ddf43f6873744
Binary files /dev/null and b/imgs/GUI_Instructions.jpg differ
diff --git a/imgs/example/SZU.jpg b/imgs/example/SZU.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..12c38a3a2d476fe4e9d48357214edde3a832f6df
Binary files /dev/null and b/imgs/example/SZU.jpg differ
diff --git a/imgs/example/SZU_summer2winter.jpg b/imgs/example/SZU_summer2winter.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..49af7fcc9e3ea8445c7b3841ff1d4e246f1693d2
Binary files /dev/null and b/imgs/example/SZU_summer2winter.jpg differ
diff --git a/imgs/example/SZU_vangogh.jpg b/imgs/example/SZU_vangogh.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..d6d7ec0a1501704f7ef41e9275c27df1cda00826
Binary files /dev/null and b/imgs/example/SZU_vangogh.jpg differ
diff --git a/imgs/example/a_dcp.png b/imgs/example/a_dcp.png
new file mode 100644
index 0000000000000000000000000000000000000000..0450dc021f4391b099d4804ca611c615b180ae47
Binary files /dev/null and b/imgs/example/a_dcp.png differ
diff --git a/imgs/example/b_dcp.png b/imgs/example/b_dcp.png
new file mode 100644
index 0000000000000000000000000000000000000000..2095c71b566444de2a3e5f91ce974cfb1a0dd0f1
Binary files /dev/null and b/imgs/example/b_dcp.png differ
diff --git a/imgs/example/face_a_clean.jpg b/imgs/example/face_a_clean.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..f00efca8faa195b6f6fe978fccb4f9e0bdab859e
Binary files /dev/null and b/imgs/example/face_a_clean.jpg differ
diff --git a/imgs/example/face_a_mosaic.jpg b/imgs/example/face_a_mosaic.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..a898e6a40fd0d78c9b4dad8c1b3e037735609b62
Binary files /dev/null and b/imgs/example/face_a_mosaic.jpg differ
diff --git a/imgs/example/face_b_clean.jpg b/imgs/example/face_b_clean.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..432285fb7f5dcfa8711b91d499d8d7d7ec008d5a
Binary files /dev/null and b/imgs/example/face_b_clean.jpg differ
diff --git a/imgs/example/face_b_mosaic.jpg b/imgs/example/face_b_mosaic.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..8aabdacd38b79c086fccbc777d6ba9404fb6ea3c
Binary files /dev/null and b/imgs/example/face_b_mosaic.jpg differ
diff --git a/imgs/example/lena.jpg b/imgs/example/lena.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..bf55198823d616f1a8b8d9508cdfd4b97e8f31db
Binary files /dev/null and b/imgs/example/lena.jpg differ
diff --git a/imgs/example/lena_add.jpg b/imgs/example/lena_add.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..871a14c08a103549ce5b7ae4d20f2bd1b052e2bf
Binary files /dev/null and b/imgs/example/lena_add.jpg differ
diff --git a/imgs/example/lena_clean.jpg b/imgs/example/lena_clean.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..bfb97107c9214301089a0f018cb0523ba6e16e79
Binary files /dev/null and b/imgs/example/lena_clean.jpg differ
diff --git a/imgs/example/youknow.png b/imgs/example/youknow.png
new file mode 100644
index 0000000000000000000000000000000000000000..9b27ebb231757fb21e4f43a788630ad9fe5f154b
Binary files /dev/null and b/imgs/example/youknow.png differ
diff --git a/imgs/example/youknow_add.png b/imgs/example/youknow_add.png
new file mode 100644
index 0000000000000000000000000000000000000000..17cb4c0ac55025b46b3d5d1a24584ca629dbcf52
Binary files /dev/null and b/imgs/example/youknow_add.png differ
diff --git a/imgs/example/youknow_clean.png b/imgs/example/youknow_clean.png
new file mode 100644
index 0000000000000000000000000000000000000000..d6dad64d7c68b3af20d4a2f34587929263e16545
Binary files /dev/null and b/imgs/example/youknow_clean.png differ
diff --git a/models/loadmodel.py b/models/loadmodel.py
index 6257da075bd0919180395ce342609882bbb8bc36..00607d7970add0e373be78279c091e31a45fcd34 100755
--- a/models/loadmodel.py
+++ b/models/loadmodel.py
@@ -3,6 +3,7 @@ from .pix2pix_model import define_G
from .pix2pixHD_model import define_G as define_G_HD
from .unet_model import UNet
from .video_model import MosaicNet
+from .videoHD_model import MosaicNet as MosaicNet_HD
def show_paramsnumber(net,netname='net'):
parameters = sum(param.numel() for param in net.parameters())
@@ -63,7 +64,10 @@ def style(opt):
return netG
def video(opt):
- netG = MosaicNet(3*25+1, 3,norm = 'batch')
+ if 'HD' in opt.model_path:
+ netG = MosaicNet_HD(3*25+1, 3, norm='instance')
+ else:
+ netG = MosaicNet(3*25+1, 3,norm = 'batch')
show_paramsnumber(netG,'netG')
netG.load_state_dict(torch.load(opt.model_path))
netG.eval()
diff --git a/models/videoHD_model.py b/models/videoHD_model.py
new file mode 100644
index 0000000000000000000000000000000000000000..9f214c50daf398b55184c759a59d64cd8dd7bb0e
--- /dev/null
+++ b/models/videoHD_model.py
@@ -0,0 +1,183 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .pix2pixHD_model import *
+
+
+class encoder_2d(nn.Module):
+ def __init__(self, input_nc, output_nc, ngf=64, n_downsampling=3, n_blocks=9, norm_layer=nn.BatchNorm2d,
+ padding_type='reflect'):
+ assert(n_blocks >= 0)
+ super(encoder_2d, self).__init__()
+ activation = nn.ReLU(True)
+
+ model = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0), norm_layer(ngf), activation]
+ ### downsample
+ for i in range(n_downsampling):
+ mult = 2**i
+ model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1),
+ norm_layer(ngf * mult * 2), activation]
+
+ self.model = nn.Sequential(*model)
+ def forward(self, input):
+ return self.model(input)
+
+class decoder_2d(nn.Module):
+ def __init__(self, input_nc, output_nc, ngf=64, n_downsampling=3, n_blocks=9, norm_layer=nn.BatchNorm2d,
+ padding_type='reflect'):
+ assert(n_blocks >= 0)
+ super(decoder_2d, self).__init__()
+ activation = nn.ReLU(True)
+
+ model = []
+
+ ### resnet blocks
+ mult = 2**n_downsampling
+ for i in range(n_blocks):
+ model += [ResnetBlock(ngf * mult, padding_type=padding_type, activation=activation, norm_layer=norm_layer)]
+
+ ### upsample
+ for i in range(n_downsampling):
+ mult = 2**(n_downsampling - i)
+ # if i%2 ==0:
+ # model += [ nn.Upsample(scale_factor = 2, mode='nearest'),
+ # nn.ReflectionPad2d(1),
+ # nn.Conv2d(ngf * mult, int(ngf * mult / 2),kernel_size=3, stride=1, padding=0),
+ # norm_layer(int(ngf * mult / 2)),
+ # nn.ReLU(True)]
+ # else:
+
+ # model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=2, padding=1, output_padding=1),
+ # norm_layer(int(ngf * mult / 2)), activation]
+
+ # model += [ nn.Upsample(scale_factor = 2, mode='nearest'),
+ # nn.ReflectionPad2d(1),
+ # nn.Conv2d(ngf * mult, int(ngf * mult / 2),kernel_size=3, stride=1, padding=0),
+ # norm_layer(int(ngf * mult / 2)),
+ # nn.ReLU(True)]
+ model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=2, padding=1, output_padding=1),
+ norm_layer(int(ngf * mult / 2)), activation]
+ model += [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()]
+ self.model = nn.Sequential(*model)
+
+ def forward(self, input):
+ return self.model(input)
+
+
+class conv_3d(nn.Module):
+ def __init__(self,inchannel,outchannel,kernel_size=3,stride=2,padding=1,norm_layer_3d=nn.BatchNorm3d,use_bias=True):
+ super(conv_3d, self).__init__()
+ self.conv = nn.Sequential(
+ nn.Conv3d(inchannel, outchannel, kernel_size=kernel_size, stride=stride, padding=padding, bias=use_bias),
+ norm_layer_3d(outchannel),
+ nn.ReLU(inplace=True),
+ )
+
+ def forward(self, x):
+ x = self.conv(x)
+ return x
+
+class conv_2d(nn.Module):
+ def __init__(self,inchannel,outchannel,kernel_size=3,stride=1,padding=1,norm_layer_2d=nn.BatchNorm2d,use_bias=True):
+ super(conv_2d, self).__init__()
+ self.conv = nn.Sequential(
+ nn.ReflectionPad2d(padding),
+ nn.Conv2d(inchannel, outchannel, kernel_size=kernel_size, stride=stride, padding=0, bias=use_bias),
+ norm_layer_2d(outchannel),
+ nn.ReLU(inplace=True),
+ )
+
+ def forward(self, x):
+ x = self.conv(x)
+ return x
+
+
+class encoder_3d(nn.Module):
+ def __init__(self,in_channel,norm_layer_2d,norm_layer_3d,use_bias):
+ super(encoder_3d, self).__init__()
+ self.inconv = conv_3d(1, 64, 7, 2, 3,norm_layer_3d,use_bias)
+ self.down1 = conv_3d(64, 128, 3, 2, 1,norm_layer_3d,use_bias)
+ self.down2 = conv_3d(128, 256, 3, 2, 1,norm_layer_3d,use_bias)
+ self.down3 = conv_3d(256, 512, 3, 2, 1,norm_layer_3d,use_bias)
+ self.down4 = conv_3d(512, 1024, 3, 1, 1,norm_layer_3d,use_bias)
+ self.pool = nn.AvgPool3d((5,1,1))
+ # self.conver2d = nn.Sequential(
+ # nn.Conv2d(256*int(in_channel/4), 256, kernel_size=3, stride=1, padding=1, bias=use_bias),
+ # norm_layer_2d(256),
+ # nn.ReLU(inplace=True),
+ # )
+
+
+ def forward(self, x):
+
+ x = x.view(x.size(0),1,x.size(1),x.size(2),x.size(3))
+ x = self.inconv(x)
+ x = self.down1(x)
+ x = self.down2(x)
+ x = self.down3(x)
+ x = self.down4(x)
+ #print(x.size())
+ x = self.pool(x)
+ #print(x.size())
+ # torch.Size([1, 1024, 16, 16])
+ # torch.Size([1, 512, 5, 16, 16])
+
+
+ x = x.view(x.size(0),x.size(1),x.size(3),x.size(4))
+
+ # x = self.conver2d(x)
+
+ return x
+
+ # def __init__(self, input_nc, output_nc, ngf=64, n_downsampling=3, n_blocks=9, norm_layer=nn.BatchNorm2d,
+ # padding_type='reflect')
+
+class ALL(nn.Module):
+ def __init__(self, in_channel, out_channel,norm_layer_2d,norm_layer_3d,use_bias):
+ super(ALL, self).__init__()
+
+ self.encoder_2d = encoder_2d(4,3,64,4,norm_layer=norm_layer_2d,padding_type='reflect')
+ self.encoder_3d = encoder_3d(in_channel,norm_layer_2d,norm_layer_3d,use_bias)
+ self.decoder_2d = decoder_2d(4,3,64,4,norm_layer=norm_layer_2d,padding_type='reflect')
+ # self.shortcut_cov = conv_2d(3,64,7,1,3,norm_layer_2d,use_bias)
+ self.merge1 = conv_2d(2048,1024,3,1,1,norm_layer_2d,use_bias)
+ # self.merge2 = nn.Sequential(
+ # conv_2d(128,64,3,1,1,norm_layer_2d,use_bias),
+ # nn.ReflectionPad2d(3),
+ # nn.Conv2d(64, out_channel, kernel_size=7, padding=0),
+ # nn.Tanh()
+ # )
+
+ def forward(self, x):
+
+ N = int((x.size()[1])/3)
+ x_2d = torch.cat((x[:,int((N-1)/2)*3:(int((N-1)/2)+1)*3,:,:], x[:,N-1:N,:,:]), 1)
+ #shortcut_2d = x[:,int((N-1)/2)*3:(int((N-1)/2)+1)*3,:,:]
+
+ x_2d = self.encoder_2d(x_2d)
+ x_3d = self.encoder_3d(x)
+ #x = x_2d + x_3d
+ x = torch.cat((x_2d,x_3d),1)
+ x = self.merge1(x)
+
+ x = self.decoder_2d(x)
+ #shortcut_2d = self.shortcut_cov(shortcut_2d)
+ #x = torch.cat((x,shortcut_2d),1)
+ #x = self.merge2(x)
+
+ return x
+
+def MosaicNet(in_channel, out_channel, norm='batch'):
+
+ if norm == 'batch':
+ # norm_layer_2d = nn.BatchNorm2d
+ # norm_layer_3d = nn.BatchNorm3d
+ norm_layer_2d = functools.partial(nn.BatchNorm2d, affine=True, track_running_stats=True)
+ norm_layer_3d = functools.partial(nn.BatchNorm3d, affine=True, track_running_stats=True)
+ use_bias = False
+ elif norm == 'instance':
+ norm_layer_2d = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)
+ norm_layer_3d = functools.partial(nn.InstanceNorm3d, affine=False, track_running_stats=False)
+ use_bias = True
+
+ return ALL(in_channel, out_channel, norm_layer_2d, norm_layer_3d, use_bias)
diff --git a/train/clean/train.py b/train/clean/train.py
index 962aeeb12cb13e1df7f0a36395932adea649f850..102d08aaa635ac21ecde069eaf2a6ca822732ae7 100644
--- a/train/clean/train.py
+++ b/train/clean/train.py
@@ -16,12 +16,17 @@ from models import pix2pix_model,pix2pixHD_model,video_model,unet_model,loadmode
from matplotlib import pyplot as plt
import torch.backends.cudnn as cudnn
+'''
+--------------------------Get options--------------------------
+'''
+
opt = Options()
opt.parser.add_argument('--N',type=int,default=25, help='')
opt.parser.add_argument('--lr',type=float,default=0.0002, help='')
opt.parser.add_argument('--beta1',type=float,default=0.5, help='')
-opt.parser.add_argument('--gan', action='store_true', help='if input it, use gan')
-opt.parser.add_argument('--l2', action='store_true', help='if input it, use L2 loss')
+opt.parser.add_argument('--gan', action='store_true', help='if specified, use gan')
+opt.parser.add_argument('--l2', action='store_true', help='if specified, use L2 loss')
+opt.parser.add_argument('--hd', action='store_true', help='if specified, use HD model')
opt.parser.add_argument('--lambda_L1',type=float,default=100, help='')
opt.parser.add_argument('--lambda_gan',type=float,default=1, help='')
opt.parser.add_argument('--finesize',type=int,default=256, help='')
@@ -36,6 +41,10 @@ opt.parser.add_argument('--startiter',type=int,default=0, help='')
opt.parser.add_argument('--continuetrain', action='store_true', help='')
opt.parser.add_argument('--savename',type=str,default='MosaicNet', help='')
+
+'''
+--------------------------Init--------------------------
+'''
opt = opt.getparse()
dir_checkpoint = os.path.join('checkpoints/',opt.savename)
util.makedirs(dir_checkpoint)
@@ -54,17 +63,18 @@ print('check dataset...')
for video in videos:
video_images = os.listdir('./dataset/'+video+'/ori')
lengths.append(len(video_images))
-#unet_128
-#resnet_9blocks
-#netG = pix2pix_model.define_G(3*N+1, 3, 128, 'resnet_6blocks', norm='instance',use_dropout=True, init_type='normal', gpu_ids=[])
-netG = videoHD_model.MosaicNet(3*N+1, 3, norm=opt.norm)
+if opt.hd:
+ netG = videoHD_model.MosaicNet(3*N+1, 3, norm=opt.norm)
+else:
+ netG = video_model.MosaicNet(3*N+1, 3, norm=opt.norm)
loadmodel.show_paramsnumber(netG,'netG')
-# netG = unet_model.UNet(3*N+1, 3)
+
if opt.gan:
- netD = pix2pixHD_model.define_D(6, 64, 3, norm=opt.norm, use_sigmoid=False, num_D=2)
- #netD = pix2pix_model.define_D(3*2+1, 64, 'pixel', norm='instance')
- #netD = pix2pix_model.define_D(3*2, 64, 'basic', norm='instance')
- #netD = pix2pix_model.define_D(3*2+1, 64, 'n_layers', n_layers_D=5, norm='instance')
+ if opt.hd:
+ netD = pix2pixHD_model.define_D(6, 64, 3, norm = opt.norm, use_sigmoid=False, num_D=2)
+ else:
+ netD = pix2pix_model.define_D(3*2, 64, 'basic', norm = opt.norm)
+ netD.train()
if opt.continuetrain:
if not os.path.isfile(os.path.join(dir_checkpoint,'last_G.pth')):
@@ -82,10 +92,11 @@ optimizer_G = torch.optim.Adam(netG.parameters(), lr=opt.lr,betas=(opt.beta1, 0.
criterion_L1 = nn.L1Loss()
criterion_L2 = nn.MSELoss()
if opt.gan:
- optimizer_D = torch.optim.Adam(netG.parameters(), lr=opt.lr,betas=(opt.beta1, 0.999))
- # criterionGAN = pix2pix_model.GANLoss(gan_mode='lsgan').cuda()
- criterionGAN = pix2pixHD_model.GANLoss(tensor=torch.cuda.FloatTensor)
- netD.train()
+ optimizer_D = torch.optim.Adam(netD.parameters(), lr=opt.lr,betas=(opt.beta1, 0.999))
+ if opt.hd:
+ criterionGAN = pix2pixHD_model.GANLoss(tensor=torch.cuda.FloatTensor)
+ else:
+ criterionGAN = pix2pix_model.GANLoss(gan_mode='lsgan').cuda()
if opt.use_gpu:
netG.cuda()
@@ -94,6 +105,9 @@ if opt.use_gpu:
criterionGAN.cuda()
cudnn.benchmark = True
+'''
+--------------------------preload data--------------------------
+'''
def loaddata():
video_index = random.randint(0,len(videos)-1)
video = videos[video_index]
@@ -136,19 +150,19 @@ def preload():
# time.sleep(0.1)
except Exception as e:
print("error:",e)
-
import threading
t = threading.Thread(target=preload,args=()) #t为新创建的线程
t.daemon = True
t.start()
-
time_start=time.time()
while load_cnt < opt.perload_num:
time.sleep(0.1)
time_end=time.time()
print('load speed:',round((time_end-time_start)/opt.perload_num,3),'s/it')
-
+'''
+--------------------------train--------------------------
+'''
util.copyfile('./train.py', os.path.join(dir_checkpoint,'train.py'))
util.copyfile('../../models/videoHD_model.py', os.path.join(dir_checkpoint,'model.py'))
netG.train()
@@ -222,7 +236,7 @@ for iter in range(opt.startiter+1,opt.maxiter):
target, pred,os.path.join(dir_checkpoint,'result_train.jpg'))
except Exception as e:
print(e)
-
+ # plot
if (iter+1)%1000 == 0:
time_end = time.time()
if opt.gan:
@@ -255,7 +269,7 @@ for iter in range(opt.startiter+1,opt.maxiter):
loss_sum = [0.,0.,0.,0.]
time_start=time.time()
-
+ # save network
if (iter+1)%opt.savefreq == 0:
if iter+1 != opt.savefreq:
os.rename(os.path.join(dir_checkpoint,'last_G.pth'),os.path.join(dir_checkpoint,str(iter+1-opt.savefreq)+'G.pth'))
@@ -274,31 +288,32 @@ for iter in range(opt.startiter+1,opt.maxiter):
print('network saved.')
#test
- netG.eval()
-
- test_names = os.listdir('./test')
- test_names.sort()
- result = np.zeros((opt.finesize*2,opt.finesize*len(test_names),3), dtype='uint8')
+ if os.path.isdir('./test'):
+ netG.eval()
+
+ test_names = os.listdir('./test')
+ test_names.sort()
+ result = np.zeros((opt.finesize*2,opt.finesize*len(test_names),3), dtype='uint8')
- for cnt,test_name in enumerate(test_names,0):
- img_names = os.listdir(os.path.join('./test',test_name,'image'))
- img_names.sort()
- inputdata = np.zeros((opt.finesize,opt.finesize,3*N+1), dtype='uint8')
- for i in range(0,N):
- img = impro.imread(os.path.join('./test',test_name,'image',img_names[i]))
- img = impro.resize(img,opt.finesize)
- inputdata[:,:,i*3:(i+1)*3] = img
+ for cnt,test_name in enumerate(test_names,0):
+ img_names = os.listdir(os.path.join('./test',test_name,'image'))
+ img_names.sort()
+ inputdata = np.zeros((opt.finesize,opt.finesize,3*N+1), dtype='uint8')
+ for i in range(0,N):
+ img = impro.imread(os.path.join('./test',test_name,'image',img_names[i]))
+ img = impro.resize(img,opt.finesize)
+ inputdata[:,:,i*3:(i+1)*3] = img
- mask = impro.imread(os.path.join('./test',test_name,'mask.png'),'gray')
- mask = impro.resize(mask,opt.finesize)
- mask = impro.mask_threshold(mask,15,128)
- inputdata[:,:,-1] = mask
- result[0:opt.finesize,opt.finesize*cnt:opt.finesize*(cnt+1),:] = inputdata[:,:,int((N-1)/2)*3:(int((N-1)/2)+1)*3]
- inputdata = data.im2tensor(inputdata,bgr2rgb=False,use_gpu=opt.use_gpu,use_transform = False,is0_1 = False)
- pred = netG(inputdata)
-
- pred = data.tensor2im(pred,rgb2bgr = False, is0_1 = False)
- result[opt.finesize:opt.finesize*2,opt.finesize*cnt:opt.finesize*(cnt+1),:] = pred
+ mask = impro.imread(os.path.join('./test',test_name,'mask.png'),'gray')
+ mask = impro.resize(mask,opt.finesize)
+ mask = impro.mask_threshold(mask,15,128)
+ inputdata[:,:,-1] = mask
+ result[0:opt.finesize,opt.finesize*cnt:opt.finesize*(cnt+1),:] = inputdata[:,:,int((N-1)/2)*3:(int((N-1)/2)+1)*3]
+ inputdata = data.im2tensor(inputdata,bgr2rgb=False,use_gpu=opt.use_gpu,use_transform = False,is0_1 = False)
+ pred = netG(inputdata)
+
+ pred = data.tensor2im(pred,rgb2bgr = False, is0_1 = False)
+ result[opt.finesize:opt.finesize*2,opt.finesize*cnt:opt.finesize*(cnt+1),:] = pred
- cv2.imwrite(os.path.join(dir_checkpoint,str(iter+1)+'_test.jpg'), result)
- netG.train()
+ cv2.imwrite(os.path.join(dir_checkpoint,str(iter+1)+'_test.jpg'), result)
+ netG.train()