refactor directory

AddMosaic.py

-import sys
-import os
-import random
-
-import numpy as np
-import cv2
-import torch
-
-from models import runmodel,loadmodel
-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
-def add_mosaic_to_image(path):
-    img = cv2.imread(path)
-    mask =runmodel.run_unet_rectim(img,net,use_gpu = opt.use_gpu)
-    mask = impro.mask_threshold(mask,opt.mask_extend,opt.mask_threshold)
-    img = mosaic.addmosaic(img,mask,opt.mosaic_size,opt.output_size,model = opt.mosaic_mod)
-    return img
-
-net = loadmodel.unet(os.path.join(opt.model_dir,opt.model_name),use_gpu = opt.use_gpu)
-
-filepaths = util.Traversal(opt.input_dir)
-
-for path in filepaths:
-    if util.is_img(path):
-        img = add_mosaic_to_image(path)
-        cv2.imwrite(os.path.join(opt.result_dir,os.path.basename(path)),img)
-    elif util.is_video(path):
-        util.clean_tempfiles()
-        fps = ffmpeg.get_video_infos(path)[0]
-        ffmpeg.video2voice(path,'./tmp/voice_tmp.mp3')
-        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)
-            img = add_mosaic_to_image(imagepath)
-            cv2.imwrite(os.path.join('./tmp/addmosaic_image',
-                                        os.path.basename(imagepath)),img)
-        ffmpeg.image2video( fps,
-                            './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'))
-

README.md

@@ -27,12 +27,12 @@ You can download pre_trained models and test video and replace the files in the
  [[百度云,提取码7thu]](https://pan.baidu.com/s/1IG4bdIiIC9PH9-oEyae5Sg) 
 
 ### Dependencies
-This code depends on numpy, scipy, opencv-python, torchvision, available via pip install.
-### AddMosaic
+This code depends on opencv-python, available via pip install.
+### Add Mosaic
 ```bash
-python3 AddMosaic.py
+python3 deepmosaic.py
 ```
-### CleanMosaic
+### Clean Mosaic
 copy the AddMosaic video from './result' to './video_or_image'
 ```bash
 python3 CleanMosaic.py

clean_cache.py

+import os
+import shutil
+
+def findalldir(rootdir):
+    dir_list = []
+    for root,dirs,files in os.walk(rootdir): 
+        for dir in dirs:
+            dir_list.append(os.path.join(root,dir))
+    return(dir_list)
+
+def Traversal(filedir):
+    file_list=[]
+    dir_list = []
+    for root,dirs,files in os.walk(filedir): 
+        for file in files:
+            file_list.append(os.path.join(root,file)) 
+        for dir in dirs:
+            dir_list.append(os.path.join(root,dir))
+            Traversal(dir)
+    return file_list,dir_list
+def is_img(path):
+    ext = os.path.splitext(path)[1]
+    ext = ext.lower()
+    if ext in ['.jpg','.png','.jpeg','.bmp']:
+        return True
+    else:
+        return False
+
+def is_video(path):
+    ext = os.path.splitext(path)[1]
+    ext = ext.lower()
+    if ext in ['.mp4','.flv','.avi','.mov','.mkv','.wmv','.rmvb']:
+        return True
+    else:
+        return False
+
+file_list,dir_list = Traversal('./')
+for file in file_list:
+    if ('tmp' in file) | ('pth' in file)|('pycache' in file) | is_video(file) | is_img(file):
+        if os.path.exists(file):
+            os.remove(file)
+            print('remove file:',file)
+
+for dir in dir_list:
+    if ('tmp'in dir)|('pycache'in dir):
+        if os.path.exists(dir):
+            shutil.rmtree(dir)
+
+            # os.rmdir(dir)
+            print('remove dir:',dir)
\ No newline at end of file

CleanMosaic.py → deepmosaic.py

+import sys
 import os
 import random
 
 import numpy as np
 import cv2
 import torch
-import scipy.signal as signal
 
 from models import runmodel,loadmodel
-from util import util,ffmpeg,data
+from util import mosaic,util,ffmpeg
 from util import image_processing as impro
-from options.cleanmosaic_options import CleanOptions
+from options import Options
 
 
-opt = CleanOptions().getparse()
+opt = Options().getparse()
+util.init(opt)
 
-def get_mosaic_position(img_origin):
-    mask =runmodel.run_unet_rectim(img_origin,net_mosaic_pos,use_gpu = opt.use_gpu)
-    mask = impro.mask_threshold(mask,10,128)
-    x,y,size,area = impro.boundingSquare(mask,Ex_mul=1.5)
-    rat = min(img_origin.shape[:2])/128.0
-    x,y,size = int(rat*x),int(rat*y),int(rat*size)
-    return x,y,size
+if opt.mode == 'add':
 
-def replace_mosaic(img_origin,img_fake,x,y,size,no_father = opt.no_feather):
-    img_fake = impro.resize(img_fake,size*2)
-
-    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)
-
-filepaths = util.Traversal(opt.input_dir)
+    net = loadmodel.unet(opt)
+    path = opt.media_path
+    if util.is_img(path):
+        print('Add Mosaic:',path)
+        img = cv2.imread(path)
+        img = runmodel.add_mosaic_to_image(img,net,opt)
+        cv2.imwrite(os.path.join(opt.result_dir,os.path.basename(path)),img)
+    elif util.is_video(path):
+        util.clean_tempfiles()
+        fps = ffmpeg.get_video_infos(path)[0]
+        ffmpeg.video2voice(path,'./tmp/voice_tmp.mp3')
+        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('Add Mosaic:',imagepath)
+            img = cv2.imread(imagepath)
+            img = runmodel.add_mosaic_to_image(img,net,opt)
+            cv2.imwrite(os.path.join('./tmp/addmosaic_image',
+                                        os.path.basename(imagepath)),img)
+        ffmpeg.image2video( fps,
+                            './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'))
 
-for path in filepaths:
+elif opt.mode == 'clean':
+    netG = loadmodel.pix2pix(opt)
+    net_mosaic_pos = loadmodel.unet_clean(opt)
+    path = opt.media_path
     if util.is_img(path):
         print('Clean Mosaic:',path)
         img_origin = cv2.imread(path)
-        x,y,size = get_mosaic_position(img_origin)
+        x,y,size = runmodel.get_mosaic_position(img_origin,net_mosaic_pos,opt)
         img_result = img_origin.copy()
         if size != 0 :
             img_mosaic = img_origin[y-size:y+size,x-size:x+size]
             img_fake=runmodel.run_pix2pix(img_mosaic,netG,use_gpu = opt.use_gpu)
-            img_result = replace_mosaic(img_origin,img_fake,x,y,size)
+            img_result = impro.replace_mosaic(img_origin,img_fake,x,y,size,opt.no_feather)
         cv2.imwrite(os.path.join(opt.result_dir,os.path.basename(path)),img_result)
 
     elif util.is_video(path):
@@ -70,12 +67,12 @@ for path in filepaths:
         for imagepath in imagepaths:
             imagepath=os.path.join('./tmp/video2image',imagepath)
             img_origin = cv2.imread(imagepath)
-            x,y,size = get_mosaic_position(img_origin)
+            x,y,size = runmodel.get_mosaic_position(img_origin,net_mosaic_pos,opt)
             positions.append([x,y,size])
             print('Find Positions:',imagepath)
         
         positions =np.array(positions)
-        for i in range(3):positions[:,i] =signal.medfilt(positions[:,i],opt.medfilt_num)
+        for i in range(3):positions[:,i] = impro.medfilt(positions[:,i],opt.medfilt_num)
 
         for i,imagepath in enumerate(imagepaths,0):
             imagepath=os.path.join('./tmp/video2image',imagepath)
@@ -84,11 +81,13 @@ for path in filepaths:
             img_result = img_origin.copy()
             if size != 0:
                 img_mosaic = img_origin[y-size:y+size,x-size:x+size]
-                img_fake=runmodel.run_pix2pix(img_mosaic,netG,use_gpu = opt.use_gpu)
-                img_result = replace_mosaic(img_origin,img_fake,x,y,size)
+                img_fake = runmodel.run_pix2pix(img_mosaic,netG,use_gpu = opt.use_gpu)
+                img_result = impro.replace_mosaic(img_origin,img_fake,x,y,size,opt.no_feather)
             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.'+opt.tempimage_type,
                     './tmp/voice_tmp.mp3',
-                     os.path.join(opt.result_dir,os.path.splitext(os.path.basename(path))[0]+'_CleanMosaic.mp4'))
+                     os.path.join(opt.result_dir,os.path.splitext(os.path.basename(path))[0]+'_CleanMosaic.mp4'))                      
+
+util.clean_tempfiles()
\ No newline at end of file

models/loadmodel.py

@@ -2,20 +2,28 @@ import torch
 from .pix2pix_model import *
 from .unet_model import UNet
 
-def pix2pix(model_path,G_model_type,use_gpu = True):
-    
-    netG = define_G(3, 3, 64, G_model_type, norm='batch',use_dropout=True, init_type='normal', gpu_ids=[])
-    netG.load_state_dict(torch.load(model_path))
+def pix2pix(opt):
+    print(opt.model_path,opt.netG)
+    netG = define_G(3, 3, 64, opt.netG, norm='batch',use_dropout=True, init_type='normal', gpu_ids=[])
+
+    netG.load_state_dict(torch.load(opt.model_path))
     netG.eval()
-    if use_gpu:
+    if opt.use_gpu:
         netG.cuda()
     return netG
 
+def unet_clean(opt):
+    net = UNet(n_channels = 3, n_classes = 1)
+    net.load_state_dict(torch.load(opt.mosaic_position_model_path))
+    net.eval()
+    if opt.use_gpu:
+        net.cuda()
+    return net
 
-def unet(model_path,use_gpu = True):
+def unet(opt):
     net = UNet(n_channels = 3, n_classes = 1)
-    net.load_state_dict(torch.load(model_path))
+    net.load_state_dict(torch.load(opt.model_path))
     net.eval()
-    if use_gpu:
+    if opt.use_gpu:
         net.cuda()
     return net

models/runmodel.py

 import sys
 sys.path.append("..")
 import util.image_processing as impro
+from util import mosaic
 from util import data
 import torch
 
@@ -30,3 +31,20 @@ def run_pix2pix(img,net,size = 128,use_gpu = True):
     img_fake = net(img)
     img_fake = data.tensor2im(img_fake)
     return img_fake
+
+
+#find mosaic position in image and add mosaic to this image
+def add_mosaic_to_image(img,net,opt):
+    mask = run_unet_rectim(img,net,use_gpu = opt.use_gpu)
+    mask = impro.mask_threshold(mask,opt.mask_extend,opt.mask_threshold)
+    img = mosaic.addmosaic(img,mask,opt.mosaic_size,opt.output_size,model = opt.mosaic_mod)
+    return img
+
+
+def get_mosaic_position(img_origin,net_mosaic_pos,opt):
+    mask = run_unet_rectim(img_origin,net_mosaic_pos,use_gpu = opt.use_gpu)
+    mask = impro.mask_threshold(mask,10,128)
+    x,y,size,area = impro.boundingSquare(mask,Ex_mul=1.5)
+    rat = min(img_origin.shape[:2])/128.0
+    x,y,size = int(rat*x),int(rat*y),int(rat*size)
+    return x,y,size
\ No newline at end of file

options/addmosaic_options.py → options.py

 import argparse
 import os
 
-class AddOptions():
+class Options():
     def __init__(self):
         self.parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
         self.initialized = False
 
     def initialize(self):
-        self.parser.add_argument('--use_gpu', action='store_true', help='if true, use gpu')
-        self.parser.add_argument('--input_dir', type=str, default='./video_or_image',help='put your videos or images here')
-        self.parser.add_argument('--result_dir', type=str, default='./result',help='result will be saved here')
-        self.parser.add_argument('--model_dir', type=str, default='./pretrained_models/AddMosaic',
-                                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')
+
+        #base
+        self.parser.add_argument('--use_gpu', action='store_true', help='if input it, use gpu')
+        self.parser.add_argument('--media_path', type=str, default='./hands_test.mp4',help='your videos or images path')
+        self.parser.add_argument('--mode', type=str, default='add',help='add or clean mosaic into your media  add | clean')
+        self.parser.add_argument('--model_path', type=str, default='./pretrained_models/add_hands_128.pth',help='pretrained model path')
+        self.parser.add_argument('--result_dir', type=str, default='./result',help='output result 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')
+
+        #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')
         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='png',help='type of temp image, png | jpg, png is better but occupy more storage space')
+        
+        #AddMosaic
+        self.parser.add_argument('--netG', type=str, default='auto',help='select model to use for netG(clean mosaic) -> auto | unet_128 | resnet_9blocks')
+        self.parser.add_argument('--mosaic_position_model_path', type=str, default='./pretrained_models/mosaic_position.pth',
+                                help='name of model use to find mosaic position')
+        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.initialized = True
 
+
     def getparse(self):
         if not self.initialized:
             self.initialize()
         self.opt = self.parser.parse_args()
+
+        if self.opt.netG == 'auto':
+            if 'unet_128' in self.opt.model_path:
+                self.opt.netG = 'unet_128'
+            elif 'resnet_9blocks' in self.opt.model_path:
+                self.opt.netG = 'resnet_9blocks'
+
         return self.opt
\ No newline at end of file

options/cleanmosaic_options.py

-import argparse
-import os
-
-class CleanOptions():
-    def __init__(self):
-        self.parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
-        self.initialized = False
-
-    def initialize(self):
-        self.parser.add_argument('--use_gpu', action='store_true', help='if true, use gpu')
-        self.parser.add_argument('--input_dir', type=str, default='./video_or_image',help='put your videos or images here')
-        self.parser.add_argument('--result_dir', type=str, default='./result',help='result will be saved here')
-        self.parser.add_argument('--model_dir', type=str, default='./pretrained_models/CleanMosaic',
-                                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.pth',
-                                help='name of model use to find mosaic position')
-        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='png',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):
-        if not self.initialized:
-            self.initialize()
-        self.opt = self.parser.parse_args()
-        return self.opt
\ No newline at end of file

pretrained_models/AddMosaic/put_AddMosaic_model_here

-put_AddMosaic_model_here

util/image_processing.py

@@ -12,6 +12,19 @@ def resize(img,size):
         res = cv2.resize(img,(size, int(size*h/w)))
     return res
 
+
+def medfilt(data,window):
+    if window%2 == 0 or window < 0:
+        print('Error: the medfilt window must be even number')
+        exit(0)
+    pad = int((window-1)/2)
+    pad_data = np.zeros(len(data)+window-1, dtype = type(data[0]))
+    result = np.zeros(len(data),dtype = type(data[0]))
+    pad_data[pad:pad+len(data)]=data[:]
+    for i in range(len(data)):
+        result[i] = np.median(pad_data[i:i+window])
+    return result
+
 def ch_one2three(img):
     #zeros = np.zeros(img.shape[:2], dtype = "uint8")
     # ret,thresh = cv2.threshold(img,127,255,cv2.THRESH_BINARY)
@@ -115,3 +128,24 @@ def mask_area(mask):
     except:
         area = 0
     return area
+
+
+def replace_mosaic(img_origin,img_fake,x,y,size,no_father):
+    img_fake = resize(img_fake,size*2)
+
+    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
\ No newline at end of file

util/util.py

@@ -25,6 +25,11 @@ def is_video(path):
     else:
         return False
 
+def init(opt):
+    if not os.path.isdir(opt.result_dir):
+        os.makedirs(opt.result_dir)
+        print('makedir:',opt.result_dir)
+        
 def  writelog(path,log):
     f = open(path,'a+')
     f.write(log+'\n')