Merge branch 'dev'

HypoX64 · May 10, 2021 · e52e6c8 · e52e6c8
2 parents 538ccbc + e100118
commit e52e6c8
Show file tree

Hide file tree

Showing 38 changed files with 1,733 additions and 1,381 deletions.
diff --git a/.gitignore b/.gitignore
@@ -143,6 +143,7 @@ video_tmp/
 result/
 nohup.out
 #./
+/.vscode
 /pix2pix
 /pix2pixHD
 /tmp
@@ -183,4 +184,5 @@ nohup.out
 *.MP4
 *.JPEG
 *.exe
-*.npy
+*.npy
+*.psd
diff --git a/README.md b/README.md
@@ -1,12 +1,14 @@
-![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).
diff --git a/README_CN.md b/README_CN.md
@@ -1,9 +1,15 @@
-![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).
 
diff --git a/cores/core.py b/cores/core.py
@@ -1,5 +1,6 @@
 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,58 +252,69 @@ 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',
                  os.path.join(opt.result_dir,os.path.splitext(os.path.basename(path))[0]+'_clean.mp4'))  
 
 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
+    N,T,S = 2,5,3
+    LEFT_FRAME = (N*S)
+    POOL_NUM = LEFT_FRAME*2+1
+    INPUT_SIZE = 256
+    FRAME_POS = np.linspace(0, (T-1)*S,T,dtype=np.int64)
+    img_pool = []
+    previous_frame = None
+    init_flag = True
+
     fps,imagepaths,height,width = video_init(opt,path)
     positions = get_mosaic_positions(opt,netM,imagepaths,savemask=True)
     t1 = time.time()
     if not opt.no_preview:
         cv2.namedWindow('clean', cv2.WINDOW_NORMAL)
 
     # clean mosaic
-    print('Clean Mosaic:')
+    print('Step:3/4 -- Clean Mosaic:')
     length = len(imagepaths)
-
-    img_pool = []
-    mosaic_input = np.zeros((INPUT_SIZE,INPUT_SIZE,3*N+1), dtype='uint8')
-
+
     for i,imagepath in enumerate(imagepaths,0):
         x,y,size = positions[i][0],positions[i][1],positions[i][2]
-
+        input_stream = []
         # image read stream
-        mask = cv2.imread(os.path.join(opt.temp_dir+'/mosaic_mask',imagepath),0)
-        if i==0 :
-            for j in range(0,N):
-                img_pool.append(impro.imread(os.path.join(opt.temp_dir+'/video2image',imagepaths[np.clip(i+j-12,0,len(imagepaths)-1)])))
-        else:
+        if i==0 :# init
+            for j in range(POOL_NUM):
+                img_pool.append(impro.imread(os.path.join(opt.temp_dir+'/video2image',imagepaths[np.clip(i+j-LEFT_FRAME,0,len(imagepaths)-1)])))
+        else: # load next frame
             img_pool.pop(0)
-            img_pool.append(impro.imread(os.path.join(opt.temp_dir+'/video2image',imagepaths[np.clip(i+12,0,len(imagepaths)-1)])))
-        img_origin = img_pool[12]
+            img_pool.append(impro.imread(os.path.join(opt.temp_dir+'/video2image',imagepaths[np.clip(i+LEFT_FRAME,0,len(imagepaths)-1)])))
+        img_origin = img_pool[LEFT_FRAME]
         img_result = img_origin.copy()
 
-        if size>100:
+        if size>50:
             try:#Avoid unknown errors
-                #reshape to network input shape
-                for k in range(N):
-                    mosaic_input[:,:,k*3:(k+1)*3] = impro.resize(img_pool[k][y-size:y+size,x-size:x+size], 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_tensor = data.im2tensor(mosaic_input,bgr2rgb=False,use_gpu=opt.use_gpu,use_transform = False,is0_1 = False)
-                unmosaic_pred = netG(mosaic_input_tensor)              
-                img_fake = data.tensor2im(unmosaic_pred,rgb2bgr = False ,is0_1 = False)
+                for pos in FRAME_POS:
+                    input_stream.append(impro.resize(img_pool[pos][y-size:y+size,x-size:x+size], INPUT_SIZE)[:,:,::-1])
+                if init_flag:
+                    init_flag = False
+                    previous_frame = input_stream[N]
+                    previous_frame = data.im2tensor(previous_frame,bgr2rgb=True,gpu_id=opt.gpu_id)
+
+                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)
+                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:
-                print('Warning:',e)
+                init_flag = True
+                print('Error:',e)
+        else:
+            init_flag = True
         cv2.imwrite(os.path.join(opt.temp_dir+'/replace_mosaic',imagepath),img_result)
         os.remove(os.path.join(opt.temp_dir+'/video2image',imagepath))      
 
@@ -301,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',