New version V0.3.0 HypoX64#5 HypoX64#8

.gitignore

@@ -154,6 +154,7 @@ result/
 /pretrained_models_old
 /deepmosaic_window
 /sftp-config.json
+/exe
 #./make_datasets
 /make_datasets/video
 /make_datasets/tmp

README.md

@@ -6,25 +6,19 @@ This porject based on "semantic segmentation" and "Image-to-Image Translation".<
 * [中文版README](./README_CN.md)<br>
 
 ### More example
-
 origin | auto add mosaic |  auto clean mosaic  
 :-:|:-:|:-:
 ![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](./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
@@ -33,6 +27,7 @@ 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 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>
@@ -64,17 +59,21 @@ You can download pre_trained models and put them into './pretrained_models'.<br>
 [[Introduction to pre-trained models]](./docs/pre-trained_models_introduction.md)<br>
 
 #### Simple example
-* Add Mosaic (output video will save in './result')<br>
+* Add Mosaic (output media 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')<br>
+* Clean Mosaic (output media 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.<br>
 [[options_introduction.md]](./docs/options_introduction.md) <br>
 
+## Training with your own dataset
+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)
+
 ## 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).
+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).
+

README_CN.md

@@ -3,25 +3,19 @@
 这是一个通过深度学习自动的为图片/视频添加马赛克,或消除马赛克的项目.<br>它基于“语义分割”以及“图像翻译”.<br>
 
 ### 更多例子
-
 原始 | 自动打码 |  自动去码  
 :-:|:-:|:-:
 ![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](./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)
 
 ## 如何运行
@@ -74,5 +68,9 @@ python3 deepmosaic.py --media_path ./result/ruoruo_add.jpg --model_path ./pretra
 如果想要测试其他的图片或视频,请参照以下文件输入参数.<br>
 [[options_introduction_CN.md]](./docs/options_introduction_CN.md) <br>
 
+## 使用自己的数据训练模型
+如果需要使用自己的数据训练模型，请参照 [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).
+代码大量的参考了以下项目:[[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).
+

cores/core.py

@@ -38,7 +38,7 @@ def addmosaic_video(opt,netS):
     positions = []
     for i,imagepath in enumerate(imagepaths,1):
         img = impro.imread(os.path.join('./tmp/video2image',imagepath))
-        mask,x,y,area = runmodel.get_ROI_position(img,netS,opt)
+        mask,x,y,size,area = runmodel.get_ROI_position(img,netS,opt)
         positions.append([x,y,area])      
         cv2.imwrite(os.path.join('./tmp/ROI_mask',imagepath),mask)
         print('\r','Find ROI location:'+str(i)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=35),end='')
@@ -110,23 +110,23 @@ 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 != 0 :
         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,x,y,size,opt.no_feather)
+        img_result = impro.replace_mosaic(img_origin,img_fake,mask,x,y,size,opt.no_feather)
     else:
         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_video_byframe(opt,netG,netM):
     path = opt.media_path
     fps,imagepaths = video_init(opt,path)[:2]
-    positions = get_mosaic_positions(opt,netM,imagepaths,savemask=False)
+    positions = get_mosaic_positions(opt,netM,imagepaths,savemask=True)
     # clean mosaic
     for i,imagepath in enumerate(imagepaths,0):
         x,y,size = positions[i][0],positions[i][1],positions[i][2]
@@ -138,7 +138,8 @@ def cleanmosaic_video_byframe(opt,netG,netM):
                 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,x,y,size,opt.no_feather)
+        mask = cv2.imread(os.path.join('./tmp/mosaic_mask',imagepath),0)
+        img_result = impro.replace_mosaic(img_origin,img_fake,mask,x,y,size,opt.no_feather)
         cv2.imwrite(os.path.join('./tmp/replace_mosaic',imagepath),img_result)
         print('\r','Clean Mosaic:'+str(i+1)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=35),end='')
     print()
@@ -178,13 +179,13 @@ def cleanmosaic_video_fusion(opt,netG,netM):
 
             mosaic_input = np.zeros((INPUT_SIZE,INPUT_SIZE,3*N+1), dtype='uint8')
             mosaic_input[:,:,0:N*3] = impro.resize(img_pool[y-size:y+size,x-size:x+size,:], INPUT_SIZE)
-            mask = impro.resize(mask,np.min(img_origin.shape[:2]))[y-size:y+size,x-size:x+size]
-            mosaic_input[:,:,-1] = impro.resize(mask, INPUT_SIZE)
+            mask_input = impro.resize(mask,np.min(img_origin.shape[:2]))[y-size:y+size,x-size:x+size]
+            mosaic_input[:,:,-1] = impro.resize(mask_input, INPUT_SIZE)
 
             mosaic_input = data.im2tensor(mosaic_input,bgr2rgb=False,use_gpu=opt.use_gpu,use_transform = False,is0_1 = False)
             unmosaic_pred = netG(mosaic_input)
             img_fake = data.tensor2im(unmosaic_pred,rgb2bgr = False ,is0_1 = False)
-            img_result = impro.replace_mosaic(img_origin,img_fake,x,y,size,opt.no_feather)
+            img_result = impro.replace_mosaic(img_origin,img_fake,mask,x,y,size,opt.no_feather)
             cv2.imwrite(os.path.join('./tmp/replace_mosaic',imagepath),img_result)
         print('\r','Clean Mosaic:'+str(i+1)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=35),end='')
     print()

cores/options.py

@@ -16,17 +16,17 @@ def initialize(self):
         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('--tempimage_type', type=str, default='jpg',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 media, if 0 -> origin')
-
+        self.parser.add_argument('--mask_threshold', type=int, default=64,help='threshold of recognize clean or add mosaic position 0~255')
+
         #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='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')

deepmosaic.py

@@ -15,7 +15,7 @@ def main():
     else:
         files = [opt.media_path]        
     if opt.mode == 'add':
-        netS = loadmodel.unet(opt)
+        netS = loadmodel.bisenet(opt,'roi')
         for file in files:
             opt.media_path = file
             if util.is_img(file):
@@ -26,7 +26,7 @@ def main():
                 print('This type of file is not supported')
 
     elif opt.mode == 'clean':
-        netM = loadmodel.unet_clean(opt)
+        netM = loadmodel.bisenet(opt,'mosaic')
         if opt.traditional:
             netG = None
         elif opt.netG == 'video':

docs/Release_notes.txt

@@ -0,0 +1,23 @@
+DeepMosaics V0.3.0
+Core program building with windows10_1703_x86_64 
+ + python 3.68  
+ + pyinstaller 3.5
+GUI building with C#
+For more detail, please view on github: https://github.com/HypoX64/DeepMosaics
+
+Releases History
+ V0.3.0
+   1. Support BiSeNet(Better recognition of mosaics).
+   2. New videoHD model.
+   3. Better feathering method.
+ V0.2.0
+   1. Add video model.
+   2. Now you can input chinese path
+   3. Support style transfer
+   4. Support fps limit
+ V0.1.2
+   1. Support pix2pixHD model
+ V0.1.1
+   1. Check path, can't input illegal path
+ V0.1.0
+   1. Initial release.