diff --git a/README.md b/README.md
index a4c6870d0e41900078942cefabf305fe1723ac1b..50460c1ec71afaf3234de8f0b5b9c21b9725e958 100755
--- a/README.md
+++ b/README.md
@@ -19,9 +19,8 @@ Download this version via [[Google Drive]](https://drive.google.com/open?id=1LTE
Attentions:
- Require Windows_x86_64, Windows10 is better.
-- File path cannot contain spaces (" ").
- Run time depends on computer performance.
-- If output video cannot be played, you can try it with [potplayer](https://daumpotplayer.com/download/).
+- If output video cannot be played, you can try with [potplayer](https://daumpotplayer.com/download/).
### Run from source
#### Prerequisites
diff --git a/clean_cache.py b/clean_cache.py
index 5bf612a8ccbdff1e510dce1165868786aa957892..a6f5104fc62c647aa1c2f270338b07702e2b4a3b 100644
--- a/clean_cache.py
+++ b/clean_cache.py
@@ -18,6 +18,7 @@ def Traversal(filedir):
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()
@@ -38,13 +39,12 @@ 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)
+ if 'imgs' not in 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
index d7202e3402cb69b4210c97739542096d291fa65b..63583cb683bb8d35775e183d54e652e9140316d1 100644
--- a/deepmosaic.py
+++ b/deepmosaic.py
@@ -7,7 +7,7 @@ import cv2
import torch
from models import runmodel,loadmodel
-from util import mosaic,util,ffmpeg
+from util import mosaic,util,ffmpeg,filt
from util import image_processing as impro
from options import Options
@@ -31,17 +31,34 @@ if opt.mode == 'add':
ffmpeg.video2image(path,'./tmp/video2image/output_%05d.'+opt.tempimage_type)
imagepaths=os.listdir('./tmp/video2image')
imagepaths.sort()
+
+ # get position
+ positions = []
for imagepath in imagepaths:
imagepath = os.path.join('./tmp/video2image',imagepath)
- print('Add Mosaic:',imagepath)
+ print('Find ROI location:',imagepath)
img = impro.imread(imagepath)
- img = runmodel.add_mosaic_to_image(img,net,opt)
+ mask,x,y,area = runmodel.get_ROI_position(img,net,opt)
+ positions.append([x,y,area])
+ cv2.imwrite(os.path.join('./tmp/ROI_mask',
+ os.path.basename(imagepath)),mask)
+ print('Optimized ROI locations...')
+ mask_index = filt.position_medfilt(np.array(positions), 7)
+
+ # add mosaic
+ print('Add mosaic to images...')
+ for i in range(len(imagepaths)):
+ mask_path = os.path.join('./tmp/ROI_mask',imagepaths[mask_index[i]])
+ mask = impro.imread(mask_path)
+ img = impro.imread(os.path.join('./tmp/video2image',imagepaths[i]))
+ img = mosaic.addmosaic(img, mask, opt)
cv2.imwrite(os.path.join('./tmp/addmosaic_image',
- os.path.basename(imagepath)),img)
+ os.path.basename(imagepaths[i])),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'))
+ os.path.join(opt.result_dir,os.path.splitext(os.path.basename(path))[0]+'_add.mp4'))
elif opt.mode == 'clean':
netG = loadmodel.pix2pix(opt)
@@ -66,16 +83,18 @@ elif opt.mode == 'clean':
positions = []
imagepaths=os.listdir('./tmp/video2image')
imagepaths.sort()
+
+ # get position
for imagepath in imagepaths:
imagepath=os.path.join('./tmp/video2image',imagepath)
img_origin = impro.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 in range(3):positions[:,i] = filt.medfilt(positions[:,i],opt.medfilt_num)
+ # clean mosaic
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]
@@ -90,6 +109,6 @@ elif opt.mode == 'clean':
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]+'_clean.mp4'))
util.clean_tempfiles(tmp_init = False)
\ No newline at end of file
diff --git a/models/loadmodel.py b/models/loadmodel.py
index 53add02d143e3e632dad71992057625568830a19..eaeb06b8a7c9a639068321ad65ce40c0ca1378a9 100755
--- a/models/loadmodel.py
+++ b/models/loadmodel.py
@@ -3,7 +3,7 @@ from .pix2pix_model import *
from .unet_model import UNet
def pix2pix(opt):
- print(opt.model_path,opt.netG)
+ # 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))
diff --git a/models/runmodel.py b/models/runmodel.py
index 095f49f5520b89338871d4f87e9b9f7a75650db9..7daa092055ad8c8a71d1471603c3b83fc1a45e9b 100755
--- a/models/runmodel.py
+++ b/models/runmodel.py
@@ -32,13 +32,11 @@ def run_pix2pix(img,net,size = 128,use_gpu = True):
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):
+def get_ROI_position(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
+ x,y,halfsize,area = impro.boundingSquare(mask, 1)
+ return mask,x,y,area
def get_mosaic_position(img_origin,net_mosaic_pos,opt):
diff --git a/options.py b/options.py
index eec80ffd53a16bd4c08a3c081215a3ee13c63a44..1ae126d4022e1a1a81b9123d6e80a677fadd3fb6 100644
--- a/options.py
+++ b/options.py
@@ -17,9 +17,9 @@ class Options():
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('--mosaic_mod', type=str, default='squa_avg',help='type of mosaic -> squa_avg | squa_random | squa_avg_circle_edge | rect_avg | random')
+ self.parser.add_argument('--mosaic_size', type=int, default=0,help='mosaic size,if 0 auto size')
+ self.parser.add_argument('--mask_extend', type=int, default=10,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')
@@ -45,6 +45,6 @@ class Options():
if self.opt.mosaic_position_model_path == 'auto':
_path = os.path.join(os.path.split(self.opt.model_path)[0],'mosaic_position.pth')
self.opt.mosaic_position_model_path = _path
- print(self.opt.mosaic_position_model_path)
+ # print(self.opt.mosaic_position_model_path)
return self.opt
\ No newline at end of file
diff --git a/util/ffmpeg.py b/util/ffmpeg.py
index b1105af440be9dfa0cf3878559cf6a4a179aa039..94b57bde937261a4932cc79bae43c42a93ad8f30 100755
--- a/util/ffmpeg.py
+++ b/util/ffmpeg.py
@@ -1,7 +1,7 @@
import os,json
def video2image(videopath,imagepath):
- os.system('ffmpeg -i '+videopath+' -f image2 '+imagepath)
+ os.system('ffmpeg -i "'+videopath+'" -f image2 '+imagepath)
def video2voice(videopath,voicepath):
os.system('ffmpeg -i '+videopath+' -f mp3 '+voicepath)
@@ -12,11 +12,11 @@ def image2video(fps,imagepath,voicepath,videopath):
os.system('ffmpeg -i ./tmp/video_tmp.mp4 -i '+voicepath+' -vcodec copy -acodec copy '+videopath)
def get_video_infos(videopath):
- cmd_str = 'ffprobe -v quiet -print_format json -show_format -show_streams -i "' + videopath + '"'
+ cmd_str = 'ffprobe -v quiet -print_format json -show_format -show_streams -i "' + videopath + '"'
out_string = os.popen(cmd_str).read()
infos = json.loads(out_string)
fps = eval(infos['streams'][0]['avg_frame_rate'])
- endtime = float(infos['streams'][0]['duration'])
+ endtime = float(infos['format']['duration'])
width = int(infos['streams'][0]['width'])
height = int(infos['streams'][0]['height'])
return fps,endtime,width,height
diff --git a/util/filt.py b/util/filt.py
new file mode 100644
index 0000000000000000000000000000000000000000..b99d6e1a100927c4ef6decd1db468b16cc69f73c
--- /dev/null
+++ b/util/filt.py
@@ -0,0 +1,77 @@
+import numpy as np
+
+def less_zero(arr,num = 7):
+ index = np.linspace(0,len(arr)-1,len(arr),dtype='int')
+ cnt = 0
+ for i in range(2,len(arr)-2):
+ if arr[i] != 0:
+ arr[i] = arr[i]
+ if cnt != 0:
+ if cnt <= num*2:
+ arr[i-cnt:round(i-cnt/2)] = arr[i-cnt-1-2]
+ arr[round(i-cnt/2):i] = arr[i+2]
+ index[i-cnt:round(i-cnt/2)] = i-cnt-1-2
+ index[round(i-cnt/2):i] = i+2
+ else:
+ arr[i-cnt:i-cnt+num] = arr[i-cnt-1-2]
+ arr[i-num:i] = arr[i+2]
+ index[i-cnt:i-cnt+num] = i-cnt-1-2
+ index[i-num:i] = i+2
+ cnt = 0
+ else:
+ cnt += 1
+ return arr,index
+
+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 position_medfilt(positions,window):
+
+ x,mask_index = less_zero(positions[:,0],window)
+ y = less_zero(positions[:,1],window)[0]
+ area = less_zero(positions[:,2],window)[0]
+ x_filt = medfilt(x, window)
+ y_filt = medfilt(y, window)
+ area_filt = medfilt(area, window)
+ cnt = 0
+ for i in range(1,len(x)):
+ if 0.850000:
+ img_mosaic = addmosaic_normal(img,mask,alpha*((area-50000)/50000+16),model = model)
+ elif 2000050000:
- img_mosaic = random_mosaic_mod(img,mask,alpha*random.uniform(16,28))
+ img_mosaic = random_mod(img,mask,alpha*random.uniform(16,28))
elif 20000