From acf9e2b2e34089b8da5f3d894436608fe2c165b0 Mon Sep 17 00:00:00 2001 From: hypox64 Date: Sun, 28 Jul 2019 23:52:12 +0800 Subject: [PATCH] refactor directory --- AddMosaic.py | 47 ---------- CleanMosaic.py | 94 ------------------- README.md | 8 +- clean_cache.py | 50 ++++++++++ deepmosaic.py | 93 ++++++++++++++++++ models/loadmodel.py | 24 +++-- models/runmodel.py | 18 ++++ options.py | 46 +++++++++ options/addmosaic_options.py | 28 ------ options/cleanmosaic_options.py | 29 ------ .../AddMosaic/put_AddMosaic_model_here | 1 - ...c_model_here => put_pretrained_model_here} | 0 result/result_save_here | 0 util/image_processing.py | 34 +++++++ util/util.py | 5 + video_or_image/put_your_input_data_here | 0 16 files changed, 266 insertions(+), 211 deletions(-) delete mode 100755 AddMosaic.py delete mode 100755 CleanMosaic.py create mode 100644 clean_cache.py create mode 100644 deepmosaic.py create mode 100644 options.py delete mode 100755 options/addmosaic_options.py delete mode 100755 options/cleanmosaic_options.py delete mode 100755 pretrained_models/AddMosaic/put_AddMosaic_model_here rename pretrained_models/{CleanMosaic/put_CleanMosaic_model_here => put_pretrained_model_here} (100%) mode change 100755 => 100644 delete mode 100755 result/result_save_here delete mode 100755 video_or_image/put_your_input_data_here diff --git a/AddMosaic.py b/AddMosaic.py deleted file mode 100755 index 2dc848b..0000000 --- a/AddMosaic.py +++ /dev/null @@ -1,47 +0,0 @@ -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')) - diff --git a/CleanMosaic.py b/CleanMosaic.py deleted file mode 100755 index 7a41bab..0000000 --- a/CleanMosaic.py +++ /dev/null @@ -1,94 +0,0 @@ -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 image_processing as impro -from options.cleanmosaic_options import CleanOptions - - -opt = CleanOptions().getparse() - -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 - -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) - -for path in filepaths: - if util.is_img(path): - print('Clean Mosaic:',path) - img_origin = cv2.imread(path) - x,y,size = get_mosaic_position(img_origin) - 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) - cv2.imwrite(os.path.join(opt.result_dir,os.path.basename(path)),img_result) - - 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) - positions = [] - imagepaths=os.listdir('./tmp/video2image') - imagepaths.sort() - for imagepath in imagepaths: - imagepath=os.path.join('./tmp/video2image',imagepath) - img_origin = cv2.imread(imagepath) - x,y,size = get_mosaic_position(img_origin) - 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,imagepath in enumerate(imagepaths,0): - imagepath=os.path.join('./tmp/video2image',imagepath) - x,y,size = positions[i][0],positions[i][1],positions[i][2] - img_origin = cv2.imread(imagepath) - 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) - 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')) diff --git a/README.md b/README.md index 01d0834..3019da7 100755 --- a/README.md +++ b/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 diff --git a/clean_cache.py b/clean_cache.py new file mode 100644 index 0000000..5bf612a --- /dev/null +++ b/clean_cache.py @@ -0,0 +1,50 @@ +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 diff --git a/deepmosaic.py b/deepmosaic.py new file mode 100644 index 0000000..aa57450 --- /dev/null +++ b/deepmosaic.py @@ -0,0 +1,93 @@ +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 import Options + + +opt = Options().getparse() +util.init(opt) + +if opt.mode == 'add': + + 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')) + +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 = 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 = 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): + 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) + positions = [] + imagepaths=os.listdir('./tmp/video2image') + imagepaths.sort() + for imagepath in imagepaths: + imagepath=os.path.join('./tmp/video2image',imagepath) + img_origin = cv2.imread(imagepath) + 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] = impro.medfilt(positions[:,i],opt.medfilt_num) + + for i,imagepath in enumerate(imagepaths,0): + imagepath=os.path.join('./tmp/video2image',imagepath) + x,y,size = positions[i][0],positions[i][1],positions[i][2] + img_origin = cv2.imread(imagepath) + 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 = 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')) + +util.clean_tempfiles() \ No newline at end of file diff --git a/models/loadmodel.py b/models/loadmodel.py index c03fb25..53add02 100755 --- a/models/loadmodel.py +++ b/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 diff --git a/models/runmodel.py b/models/runmodel.py index 3cd4567..095f49f 100755 --- a/models/runmodel.py +++ b/models/runmodel.py @@ -1,6 +1,7 @@ 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 diff --git a/options.py b/options.py new file mode 100644 index 0000000..6395579 --- /dev/null +++ b/options.py @@ -0,0 +1,46 @@ +import argparse +import os + +class Options(): + def __init__(self): + self.parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) + self.initialized = False + + def initialize(self): + + #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') + + #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 diff --git a/options/addmosaic_options.py b/options/addmosaic_options.py deleted file mode 100755 index 28c8032..0000000 --- a/options/addmosaic_options.py +++ /dev/null @@ -1,28 +0,0 @@ -import argparse -import os - -class AddOptions(): - 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') - 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') - 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 diff --git a/options/cleanmosaic_options.py b/options/cleanmosaic_options.py deleted file mode 100755 index 36edd75..0000000 --- a/options/cleanmosaic_options.py +++ /dev/null @@ -1,29 +0,0 @@ -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 diff --git a/pretrained_models/AddMosaic/put_AddMosaic_model_here b/pretrained_models/AddMosaic/put_AddMosaic_model_here deleted file mode 100755 index e7ba362..0000000 --- a/pretrained_models/AddMosaic/put_AddMosaic_model_here +++ /dev/null @@ -1 +0,0 @@ -put_AddMosaic_model_here diff --git a/pretrained_models/CleanMosaic/put_CleanMosaic_model_here b/pretrained_models/put_pretrained_model_here old mode 100755 new mode 100644 similarity index 100% rename from pretrained_models/CleanMosaic/put_CleanMosaic_model_here rename to pretrained_models/put_pretrained_model_here diff --git a/result/result_save_here b/result/result_save_here deleted file mode 100755 index e69de29..0000000 diff --git a/util/image_processing.py b/util/image_processing.py index 55ba738..890b3d5 100755 --- a/util/image_processing.py +++ b/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 diff --git a/util/util.py b/util/util.py index 111f571..894bc43 100755 --- a/util/util.py +++ b/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') diff --git a/video_or_image/put_your_input_data_here b/video_or_image/put_your_input_data_here deleted file mode 100755 index e69de29..0000000 -- GitLab