diff --git a/AddMosaic.py b/AddMosaic.py index 27bd96c721967da0941caea445f2a4166fe8fb84..2dc848bc61705ec6e59604d7e88b3e0b582df2cc 100644 --- a/AddMosaic.py +++ b/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')) diff --git a/CleanMosaic.py b/CleanMosaic.py index cdbf8a6baf3075098bde5e8d4535f384ed19af63..c710e3fc0832e16ccdf94828d45a2d6c5801f6df 100644 --- a/CleanMosaic.py +++ b/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')) diff --git a/README.md b/README.md index 8a18d016eee03afca90ed2ca18ef5ed33544452a..ce13539c4be1744b6926b83d8de902f7bf09898d 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,41 @@ ![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.
+This porject based on semantic segmentation and pix2pix. +
+## Notes +The code do not include the part of training, I will finish it in my free time. +
+## 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 +
-*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) +
+## 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 diff --git a/options/addmosaic_options.py b/options/addmosaic_options.py index d8a6e9bedaf32c6bbeef1e0ca7a447603864b8ec..c5e2cc4d2154e1c33195ade2075f95b2dbb86968 100644 --- a/options/addmosaic_options.py +++ b/options/addmosaic_options.py @@ -1,7 +1,6 @@ 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): diff --git a/options/cleanmosaic_options.py b/options/cleanmosaic_options.py index b8fde64ea50ff16b6c230d184e3957b05586db47..ec9c1ca361c6a0e3e84fa74cd42e9273dd19a45f 100644 --- a/options/cleanmosaic_options.py +++ b/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): diff --git a/util/image_processing.py b/util/image_processing.py index c30affbc40edddc0ee0fd17660cf1e649fc35722..2c805e070a94b848f275cb2345db3acc25062130 100644 --- a/util/image_processing.py +++ b/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