Add options

AddMosaic.py

@@ -11,7 +11,6 @@ from util import mosaic,util,ffmpeg
 from util import image_processing as impro
 from options.addmosaic_options import AddOptions
 
-
 opt = AddOptions().getparse()
 
 #find mosaic position in image and add mosaic to this image
@@ -34,7 +33,7 @@ for path in filepaths:
         util.clean_tempfiles()
         fps = ffmpeg.get_video_infos(path)[0]
         ffmpeg.video2voice(path,'./tmp/voice_tmp.mp3')
-        ffmpeg.video2image(path,'./tmp/video2image/output_%05d.png')
+        ffmpeg.video2image(path,'./tmp/video2image/output_%05d.'+opt.tempimage_type)
         for imagepath in os.listdir('./tmp/video2image'):
             imagepath = os.path.join('./tmp/video2image',imagepath)
             print(imagepath)
@@ -42,7 +41,7 @@ for path in filepaths:
             cv2.imwrite(os.path.join('./tmp/addmosaic_image',
                                         os.path.basename(imagepath)),img)
         ffmpeg.image2video( fps,
-                            './tmp/addmosaic_image/output_%05d.png',
+                            './tmp/addmosaic_image/output_%05d.'+opt.tempimage_type,
                             './tmp/voice_tmp.mp3',
                              os.path.join(opt.result_dir,os.path.splitext(os.path.basename(path))[0]+'_AddMosaic.mp4'))
 

CleanMosaic.py

@@ -11,6 +11,7 @@ from util import util,ffmpeg,data
 from util import image_processing as impro
 from options.cleanmosaic_options import CleanOptions
 
+
 opt = CleanOptions().getparse()
 
 def get_mosaic_position(img_origin):
@@ -21,10 +22,25 @@ def get_mosaic_position(img_origin):
     x,y,size = int(rat*x),int(rat*y),int(rat*size)
     return x,y,size
 
-def replace_mosaic(img_origin,img_fake,x,y,size):
+def replace_mosaic(img_origin,img_fake,x,y,size,no_father = opt.no_feather):
     img_fake = impro.resize(img_fake,size*2)
-    img_origin[y-size:y+size,x-size:x+size]=img_fake
-    return img_origin
+
+    if no_father:
+        img_origin[y-size:y+size,x-size:x+size]=img_fake
+        img_result = img_origin
+    else:
+        eclosion_num = int(size/5)
+        entad = int(eclosion_num/2+2)
+        mask = np.zeros(img_origin.shape, dtype='uint8')
+        mask = cv2.rectangle(mask,(x-size+entad,y-size+entad),(x+size-entad,y+size-entad),(255,255,255),-1)
+        mask = (cv2.blur(mask, (eclosion_num, eclosion_num)))
+        mask = mask/255.0
+
+        img_tmp = np.zeros(img_origin.shape)
+        img_tmp[y-size:y+size,x-size:x+size]=img_fake
+        img_result = img_origin.copy()
+        img_result = (img_origin*(1-mask)+img_tmp*mask).astype('uint8')
+    return img_result
 
 netG = loadmodel.pix2pix(os.path.join(opt.model_dir,opt.model_name),opt.model_type_netG,use_gpu = opt.use_gpu)
 net_mosaic_pos = loadmodel.unet(os.path.join(opt.model_dir,opt.mosaic_position_model_name),use_gpu = opt.use_gpu)
@@ -47,7 +63,7 @@ for path in filepaths:
         util.clean_tempfiles()
         fps = ffmpeg.get_video_infos(path)[0]
         ffmpeg.video2voice(path,'./tmp/voice_tmp.mp3')
-        ffmpeg.video2image(path,'./tmp/video2image/output_%05d.png')
+        ffmpeg.video2image(path,'./tmp/video2image/output_%05d.'+opt.tempimage_type)
         positions = []
         imagepaths=os.listdir('./tmp/video2image')
         imagepaths.sort()
@@ -59,7 +75,7 @@ for path in filepaths:
             print('Find Positions:',imagepath)
         
         positions =np.array(positions)
-        for i in range(3):positions[:,i] =signal.medfilt(positions[:,i],21)
+        for i in range(3):positions[:,i] =signal.medfilt(positions[:,i],opt.medfilt_num)
 
         for i,imagepath in enumerate(imagepaths,0):
             imagepath=os.path.join('./tmp/video2image',imagepath)
@@ -73,6 +89,6 @@ for path in filepaths:
             cv2.imwrite(os.path.join('./tmp/replace_mosaic',os.path.basename(imagepath)),img_result)
             print('Clean Mosaic:',imagepath)
         ffmpeg.image2video( fps,
-                    './tmp/replace_mosaic/output_%05d.png',
+                    './tmp/replace_mosaic/output_%05d.'+opt.tempimage_type,
                     './tmp/voice_tmp.mp3',
                      os.path.join(opt.result_dir,os.path.splitext(os.path.basename(path))[0]+'_CleanMosaic.mp4'))

README.md

 ![image](https://github.com/HypoX64/DeepMosaics/blob/master/hand.gif)
 # 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 pix2pix.
+<br>
+## Notes
+The code do not include the part of training, I will finish it in my free time.
+<br>
+## Prerequisites
+- Linux, (I didn't try this code on Windows or mac machine)
+- Python 3.5+
+- ffmpeg
+- Pytroch 0.4, (I will update to 1.0)
+- CPU or NVIDIA GPU + CUDA CuDNN
+<br>
 
-*You can use it to automatically add or remove mosaics in images or videos
-
-*I will finish this porject in a few days.
+## Getting Started
+### Clone this repo:
+```bash
+git clone https://github.com/HypoX64/DeepMosaics
+```
+### Get pre_trained models and test video
+You can download pre_trained models and test video and replace the files in the project.
+[Google Drive](https://drive.google.com/open?id=1PXt3dE9Eez2xUqpemLJutwTCC0tW-D2g)
+[百度云,提取码z8vz](https://pan.baidu.com/s/1Wi8T6PE4ExTjrHVQhv3rJA) 
+### Dependencies
+This code depends on numpy, scipy, cv2, torchvision, available via pip install.
+### AddMosaic
+```bash
+python3 AddMosaic.py
+```
+### CleanMosaic
+copy the AddMosaic video from './result' to './video_or_image'
+```bash
+python3 CleanMosaic.py
+```
+### More parameters
+[addmosaic_options](https://github.com/HypoX64/DeepMosaics/blob/master/options/addmosaic_options.py)[cleanmosaic_options](https://github.com/HypoX64/DeepMosaics/blob/master/options/cleanmosaic_options.py)
+<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).
\ No newline at end of file

options/addmosaic_options.py

 import argparse
 import os
 
-
 class AddOptions():
     def __init__(self):
         self.parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
@@ -15,11 +14,11 @@ class AddOptions():
                                 help='put pre_train model here')
         self.parser.add_argument('--model_name', type=str, default='hands_128.pth',help='name of model use to Add mosaic')
         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')
-        self.parser.add_argument('--mosaic_size', type=int, default=20,help='mosaic size')
-        self.parser.add_argument('--mask_extend', type=int, default=20,help='more mosaic')
+        self.parser.add_argument('--mosaic_size', type=int, default=30,help='mosaic size')
+        self.parser.add_argument('--mask_extend', type=int, default=20,help='more mosaic area')
         self.parser.add_argument('--mask_threshold', type=int, default=64,help='threshold of recognize mosaic position 0~255')
         self.parser.add_argument('--output_size', type=int, default=0,help='size of output file,if 0 -> origin')
-        
+        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.initialized = True
 
     def getparse(self):

options/cleanmosaic_options.py

@@ -14,9 +14,12 @@ class CleanOptions():
                                 help='put pre_train model here, including 1.model use to find mosaic position 2.model use to clean mosaic')
         self.parser.add_argument('--model_name', type=str, default='hands_unet_128.pth',help='name of model use to clean mosaic')
         self.parser.add_argument('--model_type_netG', type=str, default='unet_128',help='select model to use for netG')
-        self.parser.add_argument('--mosaic_position_model_name', type=str, default='mosaic_position.pkl',
+        self.parser.add_argument('--mosaic_position_model_name', type=str, default='mosaic_position.pth',
                                 help='name of model use to find mosaic position')
-#        self.parser.add_argument('--zoom_multiple', type=float, default=1.0,help='zoom video')
+        self.parser.add_argument('--no_feather', action='store_true', help='if true, no edge feather,but run faster')
+        self.parser.add_argument('--medfilt_num', type=int, default=11,help='medfilt window of mosaic movement in the video')
+        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('--zoom_multiple', type=float, default=1.0,help='zoom video')
         self.initialized = True
 
     def getparse(self):

util/image_processing.py

@@ -12,10 +12,10 @@ def resize(img,size):
         res = cv2.resize(img,(size, int(size*h/w)))
     return res
 
-def channel_one2three(img):
+def ch_one2three(img):
     #zeros = np.zeros(img.shape[:2], dtype = "uint8")
-    ret,thresh = cv2.threshold(img,127,255,cv2.THRESH_BINARY)
-    res = cv2.merge([thresh, thresh, thresh])
+    # ret,thresh = cv2.threshold(img,127,255,cv2.THRESH_BINARY)
+    res = cv2.merge([img, img, img])
     return res
 
 def makedataset(target_image,orgin_image):
@@ -112,4 +112,15 @@ def mask_area(mask):
         area = cv2.contourArea(contours[0])
     except:
         area = 0
-    return area
\ No newline at end of file
+    return area
+
+# def eclosion_add(img_fg,img_bg,mask,eclosion_num,Type = 'inside'):
+
+#         mask = (cv2.blur(mask, (eclosion, eclosion)))
+#         mask = ch_one2three(mask)
+#         mask = mask/255.0
+
+#         # img_tmp = np.zeros(img_bg.shape, dtype='uint8')
+#         # img_tmp[y-size:y+size,x-size:x+size]=img_fake
+#         img_result = img_origin.copy()
+#         img_result = img_origin*(1-mask)+img_tmp*mask
\ No newline at end of file