Updata README.md and fix some bugs

.gitignore

@@ -184,4 +184,5 @@ nohup.out
 *.MP4
 *.JPEG
 *.exe
-*.npy
\ No newline at end of file
+*.npy
+*.psd
\ No newline at end of file

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 project is based on "semantic segmentation" and "Image-to-Image Translation".<br>
-
-* [中文版README](./README_CN.md)<br>
+<div align="center">
+  <img src="./imgs/logo.png" width="250"><br><br>
+  <img src="https://badgen.net/github/stars/hypox64/deepmosaics?icon=github&color=4ab8a1">&emsp;<img src="https://badgen.net/github/forks/hypox64/deepmosaics?icon=github&color=4ab8a1">&emsp;<a href="https://github.com/HypoX64/DeepMosaics/releases"><img src=https://img.shields.io/github/downloads/hypox64/deepmosaics/total></a>&emsp;<a href="https://github.com/HypoX64/DeepMosaics/releases"><img src=https://img.shields.io/github/v/release/hypox64/DeepMosaics></a>&emsp;<img src=https://img.shields.io/github/license/hypox64/deepmosaics>
+</div>
 
-### More Examples
+# DeepMosaics
+**English | [中文](./README_CN.md)**<br>
+You can use it to automatically remove the mosaics in images and videos, or add mosaics to them.<br>This project is based on "semantic segmentation" and "Image-to-Image Translation".<br>Try it at this [website](http://118.89.27.46:5000/)!<br>
 
+### Examples
+![image](./imgs/hand.gif)
 origin | auto add mosaic |  auto clean mosaic  
 :-:|:-:|:-:
 ![image](./imgs/example/lena.jpg) | ![image](./imgs/example/lena_add.jpg) | ![image](./imgs/example/lena_clean.jpg) 
@@ -30,18 +32,21 @@ An interesting example:[Ricardo Milos to cat](https://www.bilibili.com/video/BV1
 ## Run DeepMosaics
 You can either run DeepMosaics via a pre-built binary package, or from source.<br>
 
+### Try it on web
+You can simply try to remove the mosaic on the face at this [website](http://118.89.27.46:5000/).<br>
 ### Pre-built binary package
 For Windows, we bulid a GUI version for easy testing.<br>
 Download this version, and a 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>
+* [[Help document]](./docs/exe_help.md)<br>
+* Video tutorial => [[youtube]](https://www.youtube.com/watch?v=1kEmYawJ_vk) [[bilibili]](https://www.bilibili.com/video/BV1QK4y1a7Av)
 
 ![image](./imgs/GUI.png)<br>
 Attentions:<br>
 
   - Requires Windows_x86_64, Windows10 is better.<br>
   - 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 computers performance(The current version does not support gpu, if you need to use gpu please run source).<br>
+  - Run time depends on computers performance (GPU version has better performance but requires CUDA to be installed).<br>
   - If output video cannot be played, you can try with [potplayer](https://daumpotplayer.com/download/).<br>
   - GUI version updates slower than source.<br>
 
@@ -67,11 +72,11 @@ You can download pre_trained models and put them into './pretrained_models'.<br>
 #### Simple Example
 * Add Mosaic (output media will save in './result')<br>
 ```bash
-python deepmosaic.py --media_path ./imgs/ruoruo.jpg --model_path ./pretrained_models/mosaic/add_face.pth --use_gpu 0
+python deepmosaic.py --media_path ./imgs/ruoruo.jpg --model_path ./pretrained_models/mosaic/add_face.pth --gpu_id 0
 ```
 * Clean Mosaic (output media will save in './result')<br>
 ```bash
-python deepmosaic.py --media_path ./result/ruoruo_add.jpg --model_path ./pretrained_models/mosaic/clean_face_HD.pth --use_gpu 0
+python deepmosaic.py --media_path ./result/ruoruo_add.jpg --model_path ./pretrained_models/mosaic/clean_face_HD.pth --gpu_id 0
 ```
 #### More Parameters
 If you want to test other images or videos, please refer to this file.<br>
@@ -81,5 +86,4 @@ If you want to test other images or videos, please refer to this file.<br>
 If you want to train with your own dataset, please refer to [training_with_your_own_dataset.md](./docs/training_with_your_own_dataset.md)
 
 ## Acknowledgements
-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) [[BiSeNet]](https://github.com/ooooverflow/BiSeNet).
-
+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) [[BiSeNet]](https://github.com/ooooverflow/BiSeNet) [[DFDNet]](https://github.com/csxmli2016/DFDNet) [[GFRNet_pytorch_new]](https://github.com/sonack/GFRNet_pytorch_new).

README_CN.md

-![image](./imgs/hand.gif)
-# <img src="./imgs/icon.jpg" width="48">DeepMosaics
-这是一个通过深度学习自动的为图片/视频添加马赛克,或消除马赛克的项目.<br>它基于“语义分割”以及“图像翻译”.<br>
+<div align="center">
+  <img src="./imgs/logo.png" width="250"><br><br>
+  <img src="https://badgen.net/github/stars/hypox64/deepmosaics?icon=github&color=4ab8a1">&emsp;<img src="https://badgen.net/github/forks/hypox64/deepmosaics?icon=github&color=4ab8a1">&emsp;<a href="https://github.com/HypoX64/DeepMosaics/releases"><img src=https://img.shields.io/github/downloads/hypox64/deepmosaics/total></a>&emsp;<a href="https://github.com/HypoX64/DeepMosaics/releases"><img src=https://img.shields.io/github/v/release/hypox64/DeepMosaics></a>&emsp;<img src=https://img.shields.io/github/license/hypox64/deepmosaics>
+</div>
+
+# DeepMosaics
+**[English](./README.md) | 中文**<br>
 
-### 更多例子
+这是一个通过深度学习自动的为图片/视频添加马赛克,或消除马赛克的项目.<br>它基于“语义分割”以及“图像翻译”.<br>现在可以在这个[网站](http://118.89.27.46:5000/)尝试使用该项目清除马赛克!<br>
 
+### 例子
+![image](./imgs/hand.gif)
 原始 | 自动打码 |  自动去码  
 :-:|:-:|:-:
 ![image](./imgs/example/lena.jpg) | ![image](./imgs/example/lena_add.jpg) | ![image](./imgs/example/lena_clean.jpg) 
@@ -26,19 +32,20 @@
 
 ## 如何运行
 可以通过我们预编译好的二进制包或源代码运行.<br>
-
+### 在网页中运行
+打开[这个网站](http://118.89.27.46:5000/)上传照片，将获得去除马赛克后的结果，受限与当地法律，目前只支持人脸.<br>
 ### 预编译的程序包
 对于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>
-
+* [[帮助文档]](./docs/exe_help_CN.md)<br>
+* [[视频教程]](https://www.bilibili.com/video/BV1QK4y1a7Av)<br>
 ![image](./imgs/GUI.png)<br>
 
 注意事项:<br>
   - 程序的运行要求在64位Windows操作系统,我仅在Windows10运行过,其他版本暂未经过测试<br>
   - 请根据需求选择合适的预训练模型进行测试，不同的预期训练模型具有不同的效果.[[预训练模型介绍]](./docs/pre-trained_models_introduction_CN.md)<br>
-  - 运行时间取决于电脑性能,对于视频文件,我们建议使用源码并在GPU上运行.<br>
+  - 运行时间取决于电脑性能,对于视频文件,我们建议在GPU上运行.<br>
   - 如果输出的视频无法播放,这边建议您尝试[potplayer](https://daumpotplayer.com/download/).<br>
   - 相比于源码,该版本的更新将会延后.
 
@@ -62,13 +69,13 @@ cd DeepMosaics
 [[预训练模型介绍]](./docs/pre-trained_models_introduction_CN.md)<br>
 
 #### 简单的例子
-* 为视频添加马赛克,例子中认为脸是需要打码的区域 ,可以通过切换预训练模型切换自动打码区域(输出结果将储存到 './result')<br>
+* 为视频或照片添加马赛克,例子中认为脸是需要打码的区域 ,可以通过切换预训练模型切换自动打码区域(输出结果将储存到 './result')<br>
 ```bash
-python deepmosaic.py --media_path ./imgs/ruoruo.jpg --model_path ./pretrained_models/mosaic/add_face.pth --use_gpu 0
+python deepmosaic.py --media_path ./imgs/ruoruo.jpg --model_path ./pretrained_models/mosaic/add_face.pth --gpu_id 0
 ```
-* 将视频中的马赛克移除,对于不同的打码物体需要使用对应的预训练模型进行马赛克消除(输出结果将储存到  './result')<br>
+* 将视频或照片中的马赛克移除,对于不同的打码物体需要使用对应的预训练模型进行马赛克消除(输出结果将储存到  './result')<br>
 ```bash
-python deepmosaic.py --media_path ./result/ruoruo_add.jpg --model_path ./pretrained_models/mosaic/clean_face_HD.pth --use_gpu 0
+python deepmosaic.py --media_path ./result/ruoruo_add.jpg --model_path ./pretrained_models/mosaic/clean_face_HD.pth --gpu_id 0
 ```
 #### 更多的参数
 如果想要测试其他的图片或视频,请参照以下文件输入参数.<br>
@@ -78,5 +85,5 @@ python deepmosaic.py --media_path ./result/ruoruo_add.jpg --model_path ./pretrai
 如果需要使用自己的数据训练模型，请参照 [training_with_your_own_dataset.md](./docs/training_with_your_own_dataset.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) [[BiSeNet]](https://github.com/ooooverflow/BiSeNet).
+代码大量的参考了以下项目:[[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) [[BiSeNet]](https://github.com/ooooverflow/BiSeNet) [[DFDNet]](https://github.com/csxmli2016/DFDNet) [[GFRNet_pytorch_new]](https://github.com/sonack/GFRNet_pytorch_new).
 

cores/core.py

 import os
 import time
+import torch
 import numpy as np
 import cv2
 
@@ -30,6 +31,7 @@ def video_init(opt,path):
                     continue_flag = True
     
     if not continue_flag:
+        print('Step:1/4 -- Convert video to images')
         util.file_init(opt)
         ffmpeg.video2voice(path,opt.temp_dir+'/voice_tmp.mp3',opt.start_time,opt.last_time)
         ffmpeg.video2image(path,opt.temp_dir+'/video2image/output_%06d.'+opt.tempimage_type,fps,opt.start_time,opt.last_time)
@@ -59,7 +61,7 @@ def addmosaic_video(opt,netS):
     if not opt.no_preview:
         cv2.namedWindow('preview', cv2.WINDOW_NORMAL)
     
-    print('Find ROI location:')
+    print('Step:2/4 -- Find ROI location')
     for i,imagepath in enumerate(imagepaths,1):
         img = impro.imread(os.path.join(opt.temp_dir+'/video2image',imagepath))
         mask,x,y,size,area = runmodel.get_ROI_position(img,netS,opt)
@@ -77,7 +79,7 @@ def addmosaic_video(opt,netS):
     mask_index = filt.position_medfilt(np.array(positions), 7)
 
     # add mosaic
-    print('Add Mosaic:')
+    print('Step:3/4 -- Add Mosaic:')
     t1 = time.time()
     for i,imagepath in enumerate(imagepaths,1):
         mask = impro.imread(os.path.join(opt.temp_dir+'/ROI_mask',imagepaths[mask_index[i-1]]),'gray')
@@ -100,6 +102,7 @@ def addmosaic_video(opt,netS):
     print()
     if not opt.no_preview:
         cv2.destroyAllWindows()
+    print('Step:4/4 -- Convert images to video')
     ffmpeg.image2video( fps,
                         opt.temp_dir+'/addmosaic_image/output_%06d.'+opt.tempimage_type,
                         opt.temp_dir+'/voice_tmp.mp3',
@@ -119,7 +122,7 @@ def styletransfer_video(opt,netG):
     path = opt.media_path
     positions = []
     fps,imagepaths = video_init(opt,path)[:2]
-    print('Transfer:')
+    print('Step:2/4 -- Transfer')
     t1 = time.time()
     if not opt.no_preview:
         cv2.namedWindow('preview', cv2.WINDOW_NORMAL)
@@ -142,6 +145,7 @@ def styletransfer_video(opt,netG):
     if not opt.no_preview:
         cv2.destroyAllWindows()
     suffix = os.path.basename(opt.model_path).replace('.pth','').replace('style_','')
+    print('Step:4/4 -- Convert images to video')
     ffmpeg.image2video( fps,
                 opt.temp_dir+'/style_transfer/output_%06d.'+opt.tempimage_type,
                 opt.temp_dir+'/voice_tmp.mp3',
@@ -156,8 +160,7 @@ def get_mosaic_positions(opt,netM,imagepaths,savemask=True):
     t1 = time.time()
     if not opt.no_preview:
         cv2.namedWindow('mosaic mask', cv2.WINDOW_NORMAL)
-
-    print('Find mosaic location:')
+    print('Step:2/4 -- Find mosaic location')
     for i,imagepath in enumerate(imagepaths,1):
         img_origin = impro.imread(os.path.join(opt.temp_dir+'/video2image',imagepath))
         x,y,size,mask = runmodel.get_mosaic_position(img_origin,netM,opt)
@@ -186,7 +189,7 @@ def cleanmosaic_img(opt,netG,netM):
     print('Clean Mosaic:',path)
     img_origin = impro.imread(path)
     x,y,size,mask = runmodel.get_mosaic_position(img_origin,netM,opt)
-    cv2.imwrite('./mask/'+os.path.basename(path), mask)
+    #cv2.imwrite('./mask/'+os.path.basename(path), mask)
     img_result = img_origin.copy()
     if size > 100 :
         img_mosaic = img_origin[y-size:y+size,x-size:x+size]
@@ -199,6 +202,18 @@ def cleanmosaic_img(opt,netG,netM):
         print('Do not find mosaic')
     impro.imwrite(os.path.join(opt.result_dir,os.path.splitext(os.path.basename(path))[0]+'_clean.jpg'),img_result)
 
+def cleanmosaic_img_server(opt,img_origin,netG,netM):
+    x,y,size,mask = runmodel.get_mosaic_position(img_origin,netM,opt)
+    img_result = img_origin.copy()
+    if size > 100 :
+        img_mosaic = img_origin[y-size:y+size,x-size:x+size]
+        if opt.traditional:
+            img_fake = runmodel.traditional_cleaner(img_mosaic,opt)
+        else:
+            img_fake = runmodel.run_pix2pix(img_mosaic,netG,opt)
+        img_result = impro.replace_mosaic(img_origin,img_fake,mask,x,y,size,opt.no_feather)
+    return img_result
+
 def cleanmosaic_video_byframe(opt,netG,netM):
     path = opt.media_path
     fps,imagepaths = video_init(opt,path)[:2]
@@ -208,7 +223,7 @@ def cleanmosaic_video_byframe(opt,netG,netM):
         cv2.namedWindow('clean', cv2.WINDOW_NORMAL)
 
     # clean mosaic
-    print('Clean Mosaic:')
+    print('Step:3/4 -- Clean Mosaic:')
     length = len(imagepaths)
     for i,imagepath in enumerate(imagepaths,0):
         x,y,size = positions[i][0],positions[i][1],positions[i][2]
@@ -237,6 +252,7 @@ def cleanmosaic_video_byframe(opt,netG,netM):
     print()
     if not opt.no_preview:
         cv2.destroyAllWindows()
+    print('Step:4/4 -- Convert images to video')
     ffmpeg.image2video( fps,
                 opt.temp_dir+'/replace_mosaic/output_%06d.'+opt.tempimage_type,
                 opt.temp_dir+'/voice_tmp.mp3',
@@ -260,7 +276,7 @@ def cleanmosaic_video_fusion(opt,netG,netM):
         cv2.namedWindow('clean', cv2.WINDOW_NORMAL)
     
     # clean mosaic
-    print('Clean Mosaic:')
+    print('Step:3/4 -- Clean Mosaic:')
     length = len(imagepaths)
     
     for i,imagepath in enumerate(imagepaths,0):
@@ -276,7 +292,7 @@ def cleanmosaic_video_fusion(opt,netG,netM):
         img_origin = img_pool[LEFT_FRAME]
         img_result = img_origin.copy()
 
-        if size>100:
+        if size>50:
             try:#Avoid unknown errors
                 for pos in FRAME_POS:
                     input_stream.append(impro.resize(img_pool[pos][y-size:y+size,x-size:x+size], INPUT_SIZE)[:,:,::-1])
@@ -287,9 +303,11 @@ def cleanmosaic_video_fusion(opt,netG,netM):
                 
                 input_stream = np.array(input_stream).reshape(1,T,INPUT_SIZE,INPUT_SIZE,3).transpose((0,4,1,2,3))
                 input_stream = data.to_tensor(data.normalize(input_stream),gpu_id=opt.gpu_id)
-                unmosaic_pred = netG(input_stream,previous_frame)
+                with torch.no_grad():
+                    unmosaic_pred = netG(input_stream,previous_frame)
                 img_fake = data.tensor2im(unmosaic_pred,rgb2bgr = True)
                 previous_frame = unmosaic_pred
+                # previous_frame = data.tensor2im(unmosaic_pred,rgb2bgr = True)
                 mask = cv2.imread(os.path.join(opt.temp_dir+'/mosaic_mask',imagepath),0)
                 img_result = impro.replace_mosaic(img_origin,img_fake,mask,x,y,size,opt.no_feather)
             except Exception as e:
@@ -309,6 +327,7 @@ def cleanmosaic_video_fusion(opt,netG,netM):
     print()
     if not opt.no_preview:
         cv2.destroyAllWindows()
+    print('Step:4/4 -- Convert images to video')
     ffmpeg.image2video( fps,
                 opt.temp_dir+'/replace_mosaic/output_%06d.'+opt.tempimage_type,
                 opt.temp_dir+'/voice_tmp.mp3',

cores/options.py

@@ -101,6 +101,8 @@ class Options():
                     self.opt.netG = 'unet_128'
                 elif 'resnet_9blocks' in model_name:
                     self.opt.netG = 'resnet_9blocks'
+                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:
@@ -116,7 +118,7 @@ class Options():
             else:
                 self.opt.ex_mult = float(self.opt.ex_mult)
 
-            if self.opt.mosaic_position_model_path == 'auto':
+            if self.opt.mosaic_position_model_path == 'auto' and self.opt.mode == 'clean':
                 _path = os.path.join(os.path.split(self.opt.model_path)[0],'mosaic_position.pth')
                 if os.path.isfile(_path):
                     self.opt.mosaic_position_model_path = _path

docs/exe_help.md

 ## DeepMosaics.exe  Instructions
-[[中文版]](./exe_help_CN.md)
+**[[中文版]](./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>
+Video tutorial => [[youtube]](https://www.youtube.com/watch?v=1kEmYawJ_vk) [[bilibili]](https://www.bilibili.com/video/BV1QK4y1a7Av)<br>
+
 Attentions:<br>
 
   - Require Windows_x86_64, Windows10 is better.<br>
@@ -9,11 +11,29 @@ Attentions:<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 install
+#### CPU version
+* 1.Download and install Microsoft Visual C++
+  https://aka.ms/vs/16/release/vc_redist.x64.exe
+#### GPU version
+Only suppport NVidia GPU above gtx1060(Driver:above 460 & CUDA:11.0)
+* 1.Download and install Microsoft Visual C++
+  https://aka.ms/vs/16/release/vc_redist.x64.exe
+* 2.Update your gpu drive to 460(or above)
+  https://www.nvidia.com/en-us/geforce/drivers/
+* 3.Download and install CUDA 11.0:
+  https://developer.nvidia.com/cuda-toolkit-archive
+
+You can also download them on BaiduNetdisk
+https://pan.baidu.com/s/10rN3U3zd5TmfGpO_PEShqQ
+Password: 1x0a
+
 ### 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'.
+* step 3: Run program and wait.
+* step 4: Cheek reult in './result'.
 
 ### Introduction to pre-trained models
 * Mosaic
@@ -22,10 +42,10 @@ Attentions:<br>
 | :------------------------------: | :---------------------------------------------------------: |
 |           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) |
+|         add_youknow.pth          |      Add mosaic to ... in images/videos.      |
+| clean_youknow_resnet_9blocks.pth |     Clean mosaic to ... in images/videos.     |
+|     clean_youknow_video.pth      |        Clean mosaic to ... in videos. It is better for processing video mosaics         |
+
 
 *  Style Transfer
 
@@ -50,8 +70,8 @@ Attentions:<br>
 * 7. More options can be input.
 * 8. Run program.
 * 9. Open help file.
-* 10.  Sponsor our project.
-* 11.  Version information.
+* 10. Sponsor our project.
+* 11. Version information.
 * 12. Open the URL on github.
 
 ### Introduction to options
@@ -60,7 +80,7 @@ If you need more effects,  use '--option your-parameters' to enter what you need
 
 |    Option    |                Description                 |                 Default                 |
 | :----------: | :----------------------------------------: | :-------------------------------------: |
-|  --use_gpu   |           if -1, do not use gpu            |                    0                    |
+|  --gpu_id   |           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 |

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>
+[视频教程](https://www.bilibili.com/video/BV1QK4y1a7Av)<br>
+
 注意事项:<br>
 
-  - 程序的运行要求在64位Windows操作系统,我仅在Windows10运行过,其他版本暂未经过测试<br>
+
+  - 程序的运行要求在64位Windows操作系统,我们仅在Windows10运行过,其他版本暂未经过测试<br>
   - 请根据需求选择合适的预训练模型进行测试<br>
-  - 运行时间取决于电脑性能,对于视频文件,我们建议使用源码以及GPU运行<br>
+  - 运行时间取决于电脑性能,对于视频文件,我们建议使用GPU运行<br>
   - 如果输出的视频无法播放,这边建议您尝试[potplayer](https://daumpotplayer.com/download/).<br>
   - 相比于源码,该版本的更新将会延后.
 
+### 如何安装
+#### CPU version
+* 1.下载安装 Microsoft Visual C++
+  https://aka.ms/vs/16/release/vc_redist.x64.exe
+#### GPU version
+仅支持gtx1060及以上的NVidia显卡(要求460版本以上的驱动以及11.0版本的CUDA, 注意只能是11.0)
+* 1.Download and install Microsoft Visual C++
+  https://aka.ms/vs/16/release/vc_redist.x64.exe
+* 2.Update your gpu drive to 460(or above)
+  https://www.nvidia.com/en-us/geforce/drivers/
+* 3.Download and install CUDA 11.0:
+  https://developer.nvidia.com/cuda-toolkit-archive
+
+当然这些也能在百度云上下载
+https://pan.baidu.com/s/10rN3U3zd5TmfGpO_PEShqQ
+提取码: 1x0a
+
 ### 如何使用
+
 * step 1: 选择需要处理的图片或视频
 * step 2: 选择预训练模型(不同的预训练模型有不同的效果)
-* step3:  运行程序并等待
-* step4:  查看结果(储存在result文件夹下)
+* step 3: 运行程序并等待
+* step 4: 查看结果(储存在result文件夹下)
 
 ## 预训练模型说明
 当前的预训练模型分为两类——添加/移除马赛克以及风格转换.
@@ -23,10 +44,10 @@
 | :------------------------------: | :-------------------------------------------: |
 |           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)  |
+|         add_youknow.pth          |        对图片或视频中的...内容打码         |
+| clean_youknow_resnet_9blocks.pth |        对图片或视频中的...内容去码         |
+|     clean_youknow_video.pth      |           对视频中的...内容去码,推荐使用带有'video'的模型去除视频中的马赛克            |
+
 
 * 风格转换
 
@@ -52,8 +73,8 @@
 * 7. 自行输入更多参数，详见下文
 * 8. 运行
 * 9. 打开帮助文件
-* 10.  支持我们
-* 11.  版本信息
+* 10. 支持我们
+* 11. 版本信息
 * 12. 打开项目的github页面
 
 ### 参数说明
@@ -62,7 +83,7 @@
 
 |    选项    |        描述         |                 默认                 |
 | :----------: | :------------------------: | :-------------------------------------: |
-|  --use_gpu   |   if -1, do not use gpu    |                    0                    |
+|  --gpu_id   |   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 |
@@ -75,7 +96,7 @@
 | --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 |
+| --mask_threshold | 马赛克区域识别阈值 0~255,越小越容易被判断为马赛克区域 | 64 |
 
 * 去除马赛克
 

docs/options_introduction.md

@@ -5,7 +5,7 @@ If you need more effects,  use '--option your-parameters' to enter what you need
 
 |    Option    |        Description         |                 Default                 |
 | :----------: | :------------------------: | :-------------------------------------: |
-|  --use_gpu   |   if -1, do not use gpu    |                    0                    |
+|  --gpu_id   |   if -1, do not use gpu    |                    0                    |
 | --media_path | your videos or images path |            ./imgs/ruoruo.jpg            |
 | --start_time | start position of video, default is the beginning of video | '00:00:00' |
 | --last_time | limit the duration of the video, default is the entire video | '00:00:00' |

docs/options_introduction_CN.md

@@ -5,7 +5,7 @@
 
 |    选项    |        描述         |                 默认                 |
 | :----------: | :------------------------: | :-------------------------------------: |
-|  --use_gpu   |   if -1, do not use gpu    |                    0                    |
+|  --gpu_id   |   if -1, do not use gpu    |                    0                    |
 | --media_path | 需要处理的视频或者照片的路径 |            ./imgs/ruoruo.jpg            |
 | --start_time | 视频开始处理的位置，默认从头开始 | '00:00:00' |
 | --last_time | 处理的视频时长，默认是整个视频 | '00:00:00' |

docs/pre-trained_models_introduction.md

@@ -10,8 +10,8 @@ Download  pre-trained model via [[Google Drive]](https://drive.google.com/open?i
 |        clean_face_HD.pth         | Clean mosaic to  faces in images/video.<br>(RAM > 8GB). |
 |         add_youknow.pth          |          Add mosaic to  ... in images/videos.           |
 | clean_youknow_resnet_9blocks.pth |         Clean mosaic to  ... in images/videos.          |
-|     clean_youknow_video.pth      |             Clean mosaic to  ... in videos.             |
-|    clean_youknow_video_HD.pth    |     Clean mosaic to  ... in videos.<br>(RAM > 8GB)      |
+|     clean_youknow_video.pth      |             Clean mosaic to  ... in videos. It is better for processing video mosaics              |
+
 
 ### Style Transfer
 

docs/pre-trained_models_introduction_CN.md

@@ -10,8 +10,8 @@
 |        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)   |
+|     clean_youknow_video.pth      |             对视频中的...内容去码,推荐使用带有'video'的模型去除视频中的马赛克               |
+
 
 ### 风格转换
 

docs/training_with_your_own_dataset.md

@@ -10,7 +10,11 @@ We will make "face" as an example. If you don't have any picture, you can downlo
   - [Pytorch 1.0+](https://pytorch.org/)
   - NVIDIA GPU(with more than 6G memory) + CUDA CuDNN<br>
 #### Dependencies
-This code depends on opencv-python, torchvision, matplotlib available via pip install.
+This code depends on opencv-python, torchvision, matplotlib, tensorboardX, scikit-image available via conda install.
+```bash
+# or
+pip install -r requirements.txt
+```
 #### Clone this repo
 ```bash
 git clone https://github.com/HypoX64/DeepMosaics
@@ -32,31 +36,31 @@ python draw_mask.py --datadir 'dir for your pictures' --savedir ../datasets/draw
 python get_image_from_video.py --datadir 'dir for your videos' --savedir ../datasets/video2image --fps 1
 ```
 ### Clean mosaic dataset
-We provide several methods for generating clean mosaic datasets. However, for better effect, we recommend train a addmosaic model in a small data  first and use it to automatically generate datasets in a big data.(recommend: Method 2(for image) & Method 4(for video))
-* Method 1: Use drawn mask to make pix2pix(HD) datasets(Require``` origin_image``` and ```mask```)
+We provide several methods for generating clean mosaic datasets. However, for better effect, we recommend train a addmosaic model in a small data  first and use it to automatically generate datasets in a big data. (recommend: Method 2(for image) & Method 4(for video))
+* Method 1: Use drawn mask to make pix2pix(HD) datasets   (Require``` origin_image``` and ```mask```)
 ```bash
 python make_pix2pix_dataset.py --datadir ../datasets/draw/face --hd --outsize 512 --fold 1 --name face --savedir ../datasets/pix2pix/face --mod drawn --minsize 128 --square
 ```
-* Method 2: Use addmosaic model to make pix2pix(HD) datasets(Require addmosaic pre-trained model)
+* Method 2: Use addmosaic model to make pix2pix(HD) datasets (Require addmosaic pre-trained model)
 ```bash
 python make_pix2pix_dataset.py --datadir 'dir for your pictures' --hd --outsize 512 --fold 1 --name face --savedir ../datasets/pix2pix/face --mod network --model_path ../pretrained_models/mosaic/add_face.pth --minsize 128 --square --mask_threshold 128
 ```
-* Method 3: Use Irregular Masks to make pix2pix(HD) datasets(Require [Irregular Masks](https://nv-adlr.github.io/publication/partialconv-inpainting))
+* Method 3: Use Irregular Masks to make pix2pix(HD) datasets (Require [Irregular Masks](https://nv-adlr.github.io/publication/partialconv-inpainting))
 ```bash
 python make_pix2pix_dataset.py --datadir 'dir for your pictures' --hd --outsize 512 --fold 1 --name face --savedir ../datasets/pix2pix/face --mod irregular --irrholedir ../datasets/Irregular_Holes_mask --square
 ```
-* Method 4: Use addmosaic model to make video datasets(Require addmosaic pre-trained model. This is better for processing video mosaics)
+* Method 4: Use addmosaic model to make video datasets (Require addmosaic pre-trained model. This is better for processing video mosaics)
 ```bash
-python make_video_dataset.py --datadir 'dir for your videos' --model_path ../pretrained_models/mosaic/add_face.pth --mask_threshold 96 --savedir ../datasets/video/face
+python make_video_dataset.py --model_path ../pretrained_models/mosaic/add_face.pth --gpu_id 0 --datadir 'dir for your videos' --savedir ../datasets/video/face
 ```
 ## Training
 ### Add
 ```bash
 cd train/add
-python train.py --use_gpu 0 --dataset ../../datasets/draw/face --savename face --loadsize 512 --finesize 360 --batchsize 16
+python train.py --gpu_id 0 --dataset ../../datasets/draw/face --savename face --loadsize 512 --finesize 360 --batchsize 16
 ```
 ### Clean
-* For image datasets(generated by ```make_pix2pix_dataset.py```)
+* For image datasets (generated by ```make_pix2pix_dataset.py```)
 We use [pix2pix](https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix) or [pix2pixHD](https://github.com/NVIDIA/pix2pixHD) to train model. We just take pix2pixHD as an example.
 ```bash
 git clone https://github.com/NVIDIA/pix2pixHD
@@ -64,10 +68,10 @@ cd pix2pixHD
 pip install dominate
 python train.py --name face --resize_or_crop resize_and_crop --loadSize 563 --fineSize 512 --label_nc 0 --no_instance --dataroot ../datasets/pix2pix/face
 ```
-* For video datasets(generated by ```make_video_dataset.py```)
+* For video datasets (generated by ```make_video_dataset.py```)
 ```bash
 cd train/clean
-python train.py --dataset ../../datasets/video/face --savename face --savefreq 100000 --gan --hd --lr 0.0002 --lambda_gan 1 --use_gpu 0
+python train.py --dataset ../../datasets/video/face --savename face --n_blocks 4 --lambda_GAN 0.01 --loadsize 286 --finesize 256 --batchsize 16 --n_layers_D 2 --num_D 3 --n_epoch 200 --gpu_id 4,5,6,7 --load_thread 16
 ```
 ## Testing
-Put saved network to ```./pretrained_models/mosaic/``` and rename it as ```add_face.pth``` or ```clean_face_HD.pth``` or ```clean_face_video_HD.pth```
+Put saved network to ```./pretrained_models/mosaic/``` and rename it as ```add_face.pth``` or ```clean_face_HD.pth``` or ```clean_face_video_HD.pth```and then run ```deepmosaic.py --model_path ./pretrained_models/mosaic/your_model_name```

make_datasets/make_pix2pix_dataset.py

@@ -16,7 +16,7 @@ import torch
 
 from models import runmodel,loadmodel
 import util.image_processing as impro
-from util import util,mosaic,data
+from util import degradater, util,mosaic,data
 
 
 opt.parser.add_argument('--datadir',type=str,default='../datasets/draw/face', help='')
@@ -107,11 +107,11 @@ for fold in range(opt.fold):
 
                 saveflag = True
                 if opt.mod == ['drawn','irregular']:
-                    x,y,size,area = impro.boundingSquare(mask_drawn, random.uniform(1.2,1.6))
+                    x,y,size,area = impro.boundingSquare(mask_drawn, random.uniform(1.1,1.6))
                 elif opt.mod == ['network','irregular']:
-                    x,y,size,area = impro.boundingSquare(mask_net, random.uniform(1.2,1.6))
+                    x,y,size,area = impro.boundingSquare(mask_net, random.uniform(1.1,1.6))
                 else:
-                    x,y,size,area = impro.boundingSquare(mask, random.uniform(1.2,1.6))
+                    x,y,size,area = impro.boundingSquare(mask, random.uniform(1.1,1.6))
 
                 if area < 1000:
                     saveflag = False
@@ -130,11 +130,15 @@ for fold in range(opt.fold):
                 if saveflag:
                     # add mosaic
                     img_mosaic = mosaic.addmosaic_random(img, mask)
-                    # random blur
+                    # random degradater
                     if random.random()>0.5:
-                        Q = random.randint(1,15)
-                        img = impro.dctblur(img,Q)
-                        img_mosaic = impro.dctblur(img_mosaic,Q)
+                        degradate_params = degradater.get_random_degenerate_params(mod='weaker_2')
+                        img = degradater.degradate(img,degradate_params)
+                        img_mosaic = degradater.degradate(img_mosaic,degradate_params)
+                    # if random.random()>0.5:
+                    #     Q = random.randint(1,15)
+                    #     img = impro.dctblur(img,Q)
+                    #     img_mosaic = impro.dctblur(img_mosaic,Q)
 
                     savecnt += 1
 

make_datasets/make_video_dataset.py

@@ -14,7 +14,7 @@ import torch
 
 from models import runmodel,loadmodel
 import util.image_processing as impro
-from util import util,mosaic,data,ffmpeg
+from util import filt, util,mosaic,data,ffmpeg
 
 
 opt.parser.add_argument('--datadir',type=str,default='your video dir', help='')
@@ -93,30 +93,65 @@ for videopath in videopaths:
             imagepaths = util.Traversal(opt.temp_dir+'/video2image')
             imagepaths = sorted(imagepaths)
             imgs=[];masks=[]
-            mask_flag = False
-
+            # mask_flag = False
+            # for imagepath in imagepaths:
+            #     img = impro.imread(imagepath)
+            #     mask = runmodel.get_ROI_position(img,net,opt,keepsize=True)[0]
+            #     imgs.append(img)
+            #     masks.append(mask)
+            #     if not mask_flag:
+            #         mask_avg = mask.astype(np.float64)
+            #         mask_flag = True
+            #     else:
+            #         mask_avg += mask.astype(np.float64)
+
+            # mask_avg = np.clip(mask_avg/len(imagepaths),0,255).astype('uint8')
+            # mask_avg = impro.mask_threshold(mask_avg,20,64)
+            # if not opt.all_mosaic_area:
+            #     mask_avg = impro.find_mostlikely_ROI(mask_avg)
+            # x,y,size,area = impro.boundingSquare(mask_avg,Ex_mul=random.uniform(1.1,1.5))
+            
+            # for i in range(len(imagepaths)):
+            #     img = impro.resize(imgs[i][y-size:y+size,x-size:x+size],opt.outsize,interpolation=cv2.INTER_CUBIC) 
+            #     mask = impro.resize(masks[i][y-size:y+size,x-size:x+size],opt.outsize,interpolation=cv2.INTER_CUBIC)
+            #     impro.imwrite(os.path.join(origindir,'%05d'%(i+1)+'.jpg'), img)
+            #     impro.imwrite(os.path.join(maskdir,'%05d'%(i+1)+'.png'), mask)
+            ex_mul = random.uniform(1.2,1.7)
+            positions = []
             for imagepath in imagepaths:
                 img = impro.imread(imagepath)
                 mask = runmodel.get_ROI_position(img,net,opt,keepsize=True)[0]
                 imgs.append(img)
                 masks.append(mask)
-                if not mask_flag:
-                    mask_avg = mask.astype(np.float64)
-                    mask_flag = True
-                else:
-                    mask_avg += mask.astype(np.float64)
-
-            mask_avg = np.clip(mask_avg/len(imagepaths),0,255).astype('uint8')
-            mask_avg = impro.mask_threshold(mask_avg,20,64)
-            if not opt.all_mosaic_area:
-                mask_avg = impro.find_mostlikely_ROI(mask_avg)
-            x,y,size,area = impro.boundingSquare(mask_avg,Ex_mul=random.uniform(1.1,1.5))
-            
-            for i in range(len(imagepaths)):
-                img = impro.resize(imgs[i][y-size:y+size,x-size:x+size],opt.outsize,interpolation=cv2.INTER_CUBIC) 
+                x,y,size,area = impro.boundingSquare(mask,Ex_mul=ex_mul)
+                positions.append([x,y,size])
+            positions =np.array(positions)
+            for i in range(3):positions[:,i] = filt.medfilt(positions[:,i],opt.medfilt_num)
+
+            for i,imagepath in enumerate(imagepaths):
+                x,y,size = positions[i][0],positions[i][1],positions[i][2]
+                tmp_cnt = i
+                while size<opt.minsize//2:
+                    tmp_cnt = tmp_cnt-1
+                    x,y,size = positions[tmp_cnt][0],positions[tmp_cnt][1],positions[tmp_cnt][2]
+                img = impro.resize(imgs[i][y-size:y+size,x-size:x+size],opt.outsize,interpolation=cv2.INTER_CUBIC)
                 mask = impro.resize(masks[i][y-size:y+size,x-size:x+size],opt.outsize,interpolation=cv2.INTER_CUBIC)
                 impro.imwrite(os.path.join(origindir,'%05d'%(i+1)+'.jpg'), img)
                 impro.imwrite(os.path.join(maskdir,'%05d'%(i+1)+'.png'), mask)
+                # x_tmp,y_tmp,size_tmp
+
+            # for imagepath in imagepaths:
+            #     img = impro.imread(imagepath)
+            #     mask,x,y,halfsize,area = runmodel.get_ROI_position(img,net,opt,keepsize=True)
+            #     if halfsize>opt.minsize//4:
+            #         if not opt.all_mosaic_area:
+            #             mask_avg = impro.find_mostlikely_ROI(mask_avg)
+            #         x,y,size,area = impro.boundingSquare(mask_avg,Ex_mul=ex_mul)
+            #     img = impro.resize(imgs[i][y-size:y+size,x-size:x+size],opt.outsize,interpolation=cv2.INTER_CUBIC)
+            #     mask = impro.resize(masks[i][y-size:y+size,x-size:x+size],opt.outsize,interpolation=cv2.INTER_CUBIC)
+            #     impro.imwrite(os.path.join(origindir,'%05d'%(i+1)+'.jpg'), img)
+            #     impro.imwrite(os.path.join(maskdir,'%05d'%(i+1)+'.png'), mask)
+
 
             result_cnt+=1
 

models/loadmodel.py

@@ -51,7 +51,7 @@ def style(opt):
     return netG
 
 def video(opt):
-    netG = video_G(N=2,n_blocks=1,gpu_id=opt.gpu_id)
+    netG = video_G(N=2,n_blocks=4,gpu_id=opt.gpu_id)
     show_paramsnumber(netG,'netG')
     netG.load_state_dict(torch.load(opt.model_path))
     netG = model_util.todevice(netG,opt.gpu_id)

requirements.txt

+opencv_python==4.5.1.48
+numpy==1.19.2
+torchvision==0.8.2
+torch==1.7.1
+matplotlib==3.3.2
+tensorboardX==2.2
+scikit-image==0.17.2
\ No newline at end of file

server.py

+import os
+import sys
+import traceback
+import cv2
+import numpy as np
+try:
+    from cores import Options,core
+    from util import util
+    from util import image_processing as impro
+    from models import loadmodel
+except Exception as e:
+    print(e)
+    input('Please press any key to exit.\n')
+    sys.exit(0)
+
+# python server.py --gpu_id 0 --model_path ./pretrained_models/mosaic/clean_face_HD.pth
+opt = Options()
+opt.parser.add_argument('--port',type=int,default=4000, help='')
+opt = opt.getparse(True)
+netM = loadmodel.bisenet(opt,'mosaic')
+netG = loadmodel.pix2pix(opt)
+
+from flask import Flask, request
+import base64
+import shutil
+
+app = Flask(__name__)
+
+@app.route("/handle", methods=["POST"])
+def handle():
+    result = {}
+    # to opencv img
+    try:
+        imgRec = request.form['img']
+        imgByte = base64.b64decode(imgRec)
+        img_np_arr = np.frombuffer(imgByte, np.uint8)
+        img = cv2.imdecode(img_np_arr, cv2.IMREAD_COLOR)
+    except Exception as e:
+        result['img'] = imgRec
+        result['info'] = 'readfailed'
+        return result
+
+    # run model
+    try:
+        if max(img.shape)>1080:
+            img = impro.resize(img,720,interpolation=cv2.INTER_CUBIC)
+        img = core.cleanmosaic_img_server(opt,img,netG,netM)
+    except Exception as e:
+        result['img'] = imgRec
+        result['info'] = 'procfailed'
+        return result
+
+    # return
+    imgbytes = cv2.imencode('.jpg', img)[1]
+    imgString = base64.b64encode(imgbytes).decode('utf-8')
+    result['img'] = imgString
+    result['info'] = 'ok'
+    return result
+
+app.run("0.0.0.0", port= opt.port, debug=opt.debug)
\ No newline at end of file

train/add/train.py

@@ -71,8 +71,8 @@ def loadimage(imagepaths,maskpaths,opt,test_flag = False):
         img,mask = data.random_transform_pair_image(img, mask, opt.finesize, test_flag)
         images[i] = (img.transpose((2, 0, 1))/255.0)
         masks[i] = (mask.reshape(1,1,opt.finesize,opt.finesize)/255.0)
-    images = Totensor(images,opt.gpu_id)
-    masks = Totensor(masks,opt.gpu_id)
+    images = data.to_tensor(images,opt.gpu_id)
+    masks = data.to_tensor(masks,opt.gpu_id)
 
     return images,masks
 

train/clean/train.py

@@ -33,13 +33,13 @@ opt.parser.add_argument('--loadsize',type=int,default=286, help='')
 opt.parser.add_argument('--batchsize',type=int,default=1, help='')
 opt.parser.add_argument('--no_gan', action='store_true', help='if specified, do not use gan')
 opt.parser.add_argument('--n_blocks',type=int,default=4, help='')
-opt.parser.add_argument('--n_layers_D',type=int,default=1, help='')
+opt.parser.add_argument('--n_layers_D',type=int,default=2, help='')
 opt.parser.add_argument('--num_D',type=int,default=3, help='')
 opt.parser.add_argument('--lambda_L2',type=float,default=100, help='')
 opt.parser.add_argument('--lambda_VGG',type=float,default=1, help='')
-opt.parser.add_argument('--lambda_GAN',type=float,default=1, help='')
+opt.parser.add_argument('--lambda_GAN',type=float,default=0.01, help='')
 opt.parser.add_argument('--lambda_D',type=float,default=1, help='')
-opt.parser.add_argument('--load_thread',type=int,default=4, help='number of thread for loading data')
+opt.parser.add_argument('--load_thread',type=int,default=16, help='number of thread for loading data')
 
 opt.parser.add_argument('--dataset',type=str,default='./datasets/face/', help='')
 opt.parser.add_argument('--dataset_test',type=str,default='./datasets/face_test/', help='')
@@ -83,6 +83,7 @@ util.makedirs(dir_checkpoint)
 localtime = time.strftime("%Y-%m-%d_%H-%M-%S", time.localtime())
 tensorboard_savedir = os.path.join('checkpoints/tensorboard',localtime+'_'+opt.savename)
 TBGlobalWriter = SummaryWriter(tensorboard_savedir)
+print('Please run "tensorboard --logdir checkpoints/tensorboardX --host=your_server_ip" and input "'+localtime+'" to filter outputs')
 
 '''
 --------------------------Init Network--------------------------

util/data.py

 import random
 import os
+from util.mosaic import get_random_parameter
 import numpy as np
 import torch
 import torchvision.transforms as transforms
@@ -134,7 +135,7 @@ def random_transform_single_image(img,finesize,params=None,test_flag = False):
             params['rate']['color'][2],params['rate']['color'][3],params['rate']['color'][4])
 
     if params['flag']['flip']:
-        img = img[:,::-1,:]
+        img = img[:,::-1]
 
     #check shape
     if img.shape[0]!= finesize or img.shape[1]!= finesize:
@@ -143,58 +144,11 @@ def random_transform_single_image(img,finesize,params=None,test_flag = False):
     return img
 
 def random_transform_pair_image(img,mask,finesize,test_flag = False):
-    #random scale
-    if random.random()<0.5:
-        h,w = img.shape[:2]
-        loadsize = min((h,w))
-        a = (float(h)/float(w))*random.uniform(0.9, 1.1)
-        if h<w:
-            mask = cv2.resize(mask, (int(loadsize/a),loadsize))
-            img = cv2.resize(img, (int(loadsize/a),loadsize))
-        else:
-            mask = cv2.resize(mask, (loadsize,int(loadsize*a)))
-            img = cv2.resize(img, (loadsize,int(loadsize*a)))
-
-    #random crop
-    h,w = img.shape[:2]
-    h_move = int((h-finesize)*random.random())
-    w_move = int((w-finesize)*random.random())
-    img_crop = img[h_move:h_move+finesize,w_move:w_move+finesize]
-    mask_crop = mask[h_move:h_move+finesize,w_move:w_move+finesize]
-
-    if test_flag:
-        return img_crop,mask_crop
-    
-    #random rotation
-    if random.random()<0.2:
-        h,w = img_crop.shape[:2]
-        M = cv2.getRotationMatrix2D((w/2,h/2),90*int(4*random.random()),1)
-        img = cv2.warpAffine(img_crop,M,(w,h))
-        mask = cv2.warpAffine(mask_crop,M,(w,h))
-    else:
-        img,mask = img_crop,mask_crop
-
-    #random color
-    img = impro.color_adjust(img,ran=True)
-
-    #random flip
-    if random.random()<0.5:
-        if random.random()<0.5:
-            img = img[:,::-1,:]
-            mask = mask[:,::-1]
-        else:
-            img = img[::-1,:,:]
-            mask = mask[::-1,:]
-
-    #random blur
-    if random.random()<0.5:
-        img = impro.dctblur(img,random.randint(1,15))
-        
-    #check shape
-    if img.shape[0]!= finesize or img.shape[1]!= finesize or mask.shape[0]!= finesize or mask.shape[1]!= finesize:
-        img = cv2.resize(img,(finesize,finesize))
-        mask = cv2.resize(mask,(finesize,finesize))
-        print('warning! shape error.')
+    params = get_transform_params()
+    img = random_transform_single_image(img,finesize,params)
+    params['flag']['degradate'] = False
+    params['flag']['color'] = False
+    mask = random_transform_single_image(mask,finesize,params)
     return img,mask
 
 def showresult(img1,img2,img3,name,is0_1 = False):

util/dataloader.py

@@ -104,7 +104,7 @@ class VideoDataLoader(object):
         self.opt = opt
         self.test_flag = test_flag
         for i in range(self.opt.n_epoch):
-            self.videolist += videolist
+            self.videolist += videolist.copy()
         random.shuffle(self.videolist)
         self.each_video_n_iter = self.opt.M -self.opt.S*(self.opt.T+1)
         self.n_iter = len(self.videolist)//self.opt.load_thread//self.opt.batchsize*self.each_video_n_iter*self.opt.load_thread