New version V0.2.0 #4 #6

.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

README.md

 ![image](./imgs/hand.gif)
 # <img src="./imgs/icon.jpg" width="48">DeepMosaics
 You can use it to automatically remove the mosaics in images and videos, or add mosaics to them.<br>
-This porject based on ‘semantic segmentation’ and ‘Image-to-Image Translation’.<br>
-Master is not stable. Please use a [stable version](https://github.com/HypoX64/DeepMosaics/tree/stable)<br>
-* [中文版](./README_CN.md)<br>
+This porject based on "semantic segmentation" and "Image-to-Image Translation".<br>
+
+* [中文版README](./README_CN.md)<br>
 
 ### More example
 origin | auto add mosaic |  auto clean mosaic  
 :-:|:-:|:-:
-![image](https://github.com/HypoX64/DeepMosaics_example/blob/master/lena.jpg) | ![image](https://github.com/HypoX64/DeepMosaics_example/blob/master/lena_add.jpg) | ![image](https://github.com/HypoX64/DeepMosaics_example/blob/master/lena_clean.jpg) 
-![image](https://github.com/HypoX64/DeepMosaics_example/blob/master/youknow.png)  | ![image](https://github.com/HypoX64/DeepMosaics_example/blob/master/youknow_add.png) | ![image](https://github.com/HypoX64/DeepMosaics_example/blob/master/youknow_clean.png) 
+![image](./imgs/example/lena.jpg) | ![image](./imgs/example/lena_add.jpg) | ![image](./imgs/example/lena_clean.jpg) 
+![image](./imgs/example/youknow.png)  | ![image](./imgs/example/youknow_add.png) | ![image](./imgs/example/youknow_clean.png) 
 * Compared with [DeepCreamPy](https://github.com/deeppomf/DeepCreamPy)
-
 mosaic image | DeepCreamPy | ours  
 :-:|:-:|:-:
-![image](https://github.com/HypoX64/DeepMosaics_example/blob/master/face_a_mosaic.jpg) | ![image](https://github.com/HypoX64/DeepMosaics_example/blob/master/a_dcp.png) | ![image](https://github.com/HypoX64/DeepMosaics_example/blob/master/face_a_clean.jpg) 
-![image](https://github.com/HypoX64/DeepMosaics_example/blob/master/face_b_mosaic.jpg) | ![image](https://github.com/HypoX64/DeepMosaics_example/blob/master/b_dcp.png) | ![image](https://github.com/HypoX64/DeepMosaics_example/blob/master/face_b_clean.jpg) 
+![image](./imgs/example/face_a_mosaic.jpg) | ![image](./imgs/example/a_dcp.png) | ![image](./imgs/example/face_a_clean.jpg) 
+![image](./imgs/example/face_b_mosaic.jpg) | ![image](./imgs/example/b_dcp.png) | ![image](./imgs/example/face_b_clean.jpg) 
+* Style Transfer
+origin | to Van Gogh | to winter
+:-:|:-:|:-:
+![image](./imgs/example/SZU.jpg) | ![image](./imgs/example/SZU_vangogh.jpg) | ![image](./imgs/example/SZU_summer2winter.jpg) 
+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.<br>
 
 ### Pre-built binary package
 For windows, we bulid a GUI version for easy test.<br>
-Download this version via [[Google Drive]](https://drive.google.com/open?id=1LTERcN33McoiztYEwBxMuRjjgxh4DEPs)  [[百度云,提取码1x0a]](https://pan.baidu.com/s/10rN3U3zd5TmfGpO_PEShqQ) <br>
+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) <br>
+* [[How to use]](./docs/exe_help.md)<br>
 
 ![image](./imgs/GUI.png)<br>
-
 Attentions:<br>
+
   - Require Windows_x86_64, Windows10 is better.<br>
-  - Different pre-trained models are suitable for different effects.<br>
-  - Run time depends on computer performance.<br>
-  - 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)<br>
+  - Run time depends on computer performance(The current version does not support gpu, if you need to use gpu please run source).<br>
+  - If output video cannot be played, you can try with [potplayer](https://daumpotplayer.com/download/).<br>
+  - GUI version update slower than source.<br>
 
 ### 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<br>
 #### 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'.<br>
-[[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)<br>
+[[Introduction to pre-trained models]](./docs/pre-trained_models_introduction.md)<br>
 
 #### Simple example
-* Add Mosaic (output video will save in './result')
+* Add Mosaic (output video will save in './result')<br>
 ```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')<br>
 ```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) <br>
+If you want to test other image or video, please refer to this file.<br>
+[[options_introduction.md]](./docs/options_introduction.md) <br>
 
 ## 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).

README_CN.md

 ![image](./imgs/hand.gif)
 # <img src="./imgs/icon.jpg" width="48">DeepMosaics
 这是一个通过深度学习自动的为图片/视频添加马赛克,或消除马赛克的项目.<br>它基于“语义分割”以及“图像翻译”.<br>
-主分支并不稳定,请移步[稳定版本](https://github.com/HypoX64/DeepMosaics/tree/stable)<br>
+
 ### 更多例子
 原始 | 自动打码 |  自动去码  
 :-:|:-:|:-:
-![image](https://github.com/HypoX64/DeepMosaics_example/blob/master/lena.jpg) | ![image](https://github.com/HypoX64/DeepMosaics_example/blob/master/lena_add.jpg) | ![image](https://github.com/HypoX64/DeepMosaics_example/blob/master/lena_clean.jpg) 
-![image](https://github.com/HypoX64/DeepMosaics_example/blob/master/youknow.png)  | ![image](https://github.com/HypoX64/DeepMosaics_example/blob/master/youknow_add.png) | ![image](https://github.com/HypoX64/DeepMosaics_example/blob/master/youknow_clean.png) 
+![image](./imgs/example/lena.jpg) | ![image](./imgs/example/lena_add.jpg) | ![image](./imgs/example/lena_clean.jpg) 
+![image](./imgs/example/youknow.png)  | ![image](./imgs/example/youknow_add.png) | ![image](./imgs/example/youknow_clean.png) 
 * 与 [DeepCreamPy](https://github.com/deeppomf/DeepCreamPy)相比较
-
 马赛克图片 | DeepCreamPy | ours  
 :-:|:-:|:-:
-![image](https://github.com/HypoX64/DeepMosaics_example/blob/master/face_a_mosaic.jpg) | ![image](https://github.com/HypoX64/DeepMosaics_example/blob/master/a_dcp.png) | ![image](https://github.com/HypoX64/DeepMosaics_example/blob/master/face_a_clean.jpg) 
-![image](https://github.com/HypoX64/DeepMosaics_example/blob/master/face_b_mosaic.jpg) | ![image](https://github.com/HypoX64/DeepMosaics_example/blob/master/b_dcp.png) | ![image](https://github.com/HypoX64/DeepMosaics_example/blob/master/face_b_clean.jpg) 
-
-## 一些说明
-* 训练部分并不完全<br>
-* 现在,代码已经支持基于[pix2pixHD](https://github.com/NVIDIA/pix2pixHD)训练出的模型,但网络仍在训练中,这将使得输出结果看起来更加清晰,"真实".<br>
-* 新的模型,可根据视频帧间关系进行马赛克恢复,在pretrained model 中被命名为*_video.pth<br>
+![image](./imgs/example/face_a_mosaic.jpg) | ![image](./imgs/example/a_dcp.png) | ![image](./imgs/example/face_a_clean.jpg) 
+![image](./imgs/example/face_b_mosaic.jpg) | ![image](./imgs/example/b_dcp.png) | ![image](./imgs/example/face_b_clean.jpg) 
+* 风格转换
+原始 | 梵高风格 | 转化为冬天
+:-:|:-:|:-:
+![image](./imgs/example/SZU.jpg) | ![image](./imgs/example/SZU_vangogh.jpg) | ![image](./imgs/example/SZU_summer2winter.jpg) 
+一个有意思的尝试:[香蕉君♂猫](https://www.bilibili.com/video/BV1Q7411W7n6)
 
 ## 如何运行
 可以通过我们预编译好的二进制包或源代码运行.<br>
@@ -26,12 +25,14 @@
 对于Windows用户,我们提供了包含GUI界面的免安装软件包.<br>
 可以通过下面两种方式进行下载: [[Google Drive]](https://drive.google.com/open?id=1LTERcN33McoiztYEwBxMuRjjgxh4DEPs)  [[百度云,提取码1x0a]](https://pan.baidu.com/s/10rN3U3zd5TmfGpO_PEShqQ) <br>
 
+* [[使用教程]](./docs/exe_help_CN.md)<br>
+
 ![image](./imgs/GUI.png)<br>
 
 注意事项:<br>
   - 程序的运行要求在64位Windows操作系统,我仅在Windows10运行过,其他版本暂未经过测试<br>
-  - 请根据需求选择合适的预训练模型进行测试<br>
-  - 运行时间取决于电脑性能,对于视频文件,我们建议可以先使用截图进行测试.<br>
+  - 请根据需求选择合适的预训练模型进行测试，不同的预期训练模型具有不同的效果.[[预训练模型介绍]](./docs/pre-trained_models_introduction_CN.md)<br>
+  - 运行时间取决于电脑性能,对于视频文件,我们建议使用源码并在GPU上运行.<br>
   - 如果输出的视频无法播放,这边建议您尝试[potplayer](https://daumpotplayer.com/download/).<br>
   - 相比于源码,该版本的更新将会延后.
 
@@ -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<br>
 #### Python依赖项
 代码依赖于opencv-python以及 torchvision,可有通过pip install 进行安装.
@@ -52,19 +53,20 @@ cd DeepMosaics
 #### 下载预训练模型
 可以通过以下两种方法下载预训练模型,并将他们置于'./pretrained_models'文件夹中.<br>
 [[Google Drive]](https://drive.google.com/open?id=1LTERcN33McoiztYEwBxMuRjjgxh4DEPs)  [[百度云,提取码1x0a]](https://pan.baidu.com/s/10rN3U3zd5TmfGpO_PEShqQ) <br>
+[[预训练模型介绍]](./docs/pre-trained_models_introduction_CN.md)<br>
 
 #### 简单的例子
-* 为视频添加马赛克,例子中认为脸是需要打码的区域 ,可以通过切换预训练模型切换自动打码区域(输出结果将储存到 './result')
+* 为视频添加马赛克,例子中认为脸是需要打码的区域 ,可以通过切换预训练模型切换自动打码区域(输出结果将储存到 './result')<br>
 ```bash
 python3 deepmosaic.py --media_path ./imgs/ruoruo.jpg --model_path ./pretrained_models/mosaic/add_face.pth --use_gpu -1
 ```
-* 将视频中的马赛克移除,对于不同的打码物体需要使用对应的预训练模型进行马赛克消除(输出结果将储存到  './result')
+* 将视频中的马赛克移除,对于不同的打码物体需要使用对应的预训练模型进行马赛克消除(输出结果将储存到  './result')<br>
 ```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) <br>
+如果想要测试其他的图片或视频,请参照以下文件输入参数.<br>
+[[options_introduction_CN.md]](./docs/options_introduction_CN.md) <br>
 
 ## 鸣谢
-代码大量的参考了以下项目:[[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

cores/core.py

@@ -150,6 +150,9 @@ def cleanmosaic_video_byframe(opt,netG,netM):
 def cleanmosaic_video_fusion(opt,netG,netM):
     path = opt.media_path
     N = 25
+    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)

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:

deepmosaic.py

 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)
+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)
+

docs/exe_help.md

+## DeepMosaics.exe  Instructions
+[[中文版]](./exe_help_CN.md)
+This is a GUI version compiled in Windows.<br>
+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) <br>
+Attentions:<br>
+
+  - Require Windows_x86_64, Windows10 is better.<br>
+  - Different pre-trained models are suitable for different effects.<br>
+  - Run time depends on computer performance.<br>
+  - If output video cannot be played, you can try with [potplayer](https://daumpotplayer.com/download/).<br>
+  - GUI version update slower than source.<br>
+### 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.<br>(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.<br>(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
+![image](../imgs/GUI_Instructions.jpg)<br>
+* 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

docs/exe_help_CN.md

+## DeepMosaics.exe  使用说明
+下载程序以及预训练模型 [[Google Drive]](https://drive.google.com/open?id=1LTERcN33McoiztYEwBxMuRjjgxh4DEPs)  [[百度云,提取码1x0a]](https://pan.baidu.com/s/10rN3U3zd5TmfGpO_PEShqQ) <br>
+注意事项:<br>
+
+  - 程序的运行要求在64位Windows操作系统,我仅在Windows10运行过,其他版本暂未经过测试<br>
+  - 请根据需求选择合适的预训练模型进行测试<br>
+  - 运行时间取决于电脑性能,对于视频文件,我们建议使用源码以及GPU运行<br>
+  - 如果输出的视频无法播放,这边建议您尝试[potplayer](https://daumpotplayer.com/download/).<br>
+  - 相比于源码,该版本的更新将会延后.
+
+### 如何使用
+* step 1: 选择需要处理的图片或视频
+* step 2: 选择预训练模型(不同的预训练模型有不同的效果)
+* step3:  运行程序并等待
+* step4:  查看结果(储存在result文件夹下)
+
+## 预训练模型说明
+当前的预训练模型分为两类——添加/移除马赛克以及风格转换.
+
+* 马赛克
+|              文件名              |                     描述                      |
+| :------------------------------: | :-------------------------------------------: |
+|           add_face.pth           |           对图片或视频中的脸部打码            |
+|        clean_face_HD.pth         | 对图片或视频中的脸部去码<br>(要求内存 > 8GB). |
+|         add_youknow.pth          |        对图片或视频中的十八禁内容打码         |
+| clean_youknow_resnet_9blocks.pth |        对图片或视频中的十八禁内容去码         |
+|     clean_youknow_video.pth      |           对视频中的十八禁内容去码            |
+|    clean_youknow_video_HD.pth    | 对视频中的十八禁内容去码<br>(要求内存 > 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界面注释
+![image](../imgs/GUI_Instructions.jpg)<br>
+* 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

docs/options_introduction.md

+## 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

docs/options_introduction_CN.md

+## 参数说明
+如果需要更多的效果,  请按照 '--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

docs/pre-trained_models_introduction.md

+## 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) <br>
+
+### 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.<br>(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.<br>(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. |
+

docs/pre-trained_models_introduction_CN.md

+## 预训练模型说明
+当前的预训练模型分为两类——添加/移除马赛克以及风格转换.
+可以通过以下方式下载预训练模型 [[Google Drive]](https://drive.google.com/open?id=1LTERcN33McoiztYEwBxMuRjjgxh4DEPs)  [[百度云,提取码1x0a]](https://pan.baidu.com/s/10rN3U3zd5TmfGpO_PEShqQ) <br>
+
+### 马赛克
+|              文件名              |                     描述                      |
+| :------------------------------: | :-------------------------------------------: |
+|           add_face.pth           |           对图片或视频中的脸部打码            |
+|        clean_face_HD.pth         | 对图片或视频中的脸部去码<br>(要求内存 > 8GB). |
+|         add_youknow.pth          |        对图片或视频中的十八禁内容打码         |
+| clean_youknow_resnet_9blocks.pth |        对图片或视频中的十八禁内容去码         |
+|     clean_youknow_video.pth      |           对视频中的十八禁内容去码            |
+|    clean_youknow_video_HD.pth    | 对视频中的十八禁内容去码<br>(要求内存 > 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的绘画风格 |
+

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,6 +64,9 @@ def style(opt):
     return netG
 
 def video(opt):
+    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))

models/videoHD_model.py

+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)

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()
+    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)
-    netD.train()
+    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,6 +288,7 @@ for iter in range(opt.startiter+1,opt.maxiter):
         print('network saved.')
 
         #test
+        if os.path.isdir('./test'):  
             netG.eval()
             
             test_names = os.listdir('./test')