diff --git a/cores/core.py b/cores/core.py index fea1dccc7c647e358dc0291f16ee04955f74dea6..c95d14766c2ff87e112d85fa746561fba85485f0 100644 --- a/cores/core.py +++ b/cores/core.py @@ -11,14 +11,32 @@ from util import image_processing as impro ---------------------Video Init--------------------- ''' def video_init(opt,path): - util.clean_tempfiles(opt) fps,endtime,height,width = ffmpeg.get_video_infos(path) if opt.fps !=0: fps = opt.fps - ffmpeg.video2voice(path,opt.temp_dir+'/voice_tmp.mp3',opt.start_time,opt.last_time) - ffmpeg.video2image(path,opt.temp_dir+'/video2image/output_%06d.'+opt.tempimage_type,fps,opt.start_time,opt.last_time) - imagepaths=os.listdir(opt.temp_dir+'/video2image') - imagepaths.sort() + + continue_flag = False + if os.path.isdir(opt.temp_dir): + if (opt.last_time != '00:00:00' and len(os.listdir(os.path.join(opt.temp_dir,'video2image'))) > fps*(util.stamp2second(opt.last_time)-1)) \ + or (opt.last_time == '00:00:00' and len(os.listdir(os.path.join(opt.temp_dir,'video2image'))) > fps*(endtime-1)): + choose = input('There is an unfinished video. Continue it? [y/n] ') + if choose.lower() =='yes' or choose.lower() == 'y': + continue_flag = True + + if continue_flag: + processed_num = max(len(os.listdir(os.path.join(opt.temp_dir,'addmosaic_image'))), + len(os.listdir(os.path.join(opt.temp_dir,'replace_mosaic'))), + len(os.listdir(os.path.join(opt.temp_dir,'style_transfer')))) + imagepaths = os.listdir(opt.temp_dir+'/video2image') + imagepaths.sort() + imagepaths = imagepaths[processed_num:] + else: + util.file_init(opt) + ffmpeg.video2voice(path,opt.temp_dir+'/voice_tmp.mp3',opt.start_time,opt.last_time) + ffmpeg.video2image(path,opt.temp_dir+'/video2image/output_%06d.'+opt.tempimage_type,fps,opt.start_time,opt.last_time) + imagepaths = os.listdir(opt.temp_dir+'/video2image') + imagepaths.sort() + return fps,imagepaths,height,width ''' @@ -238,34 +256,35 @@ def cleanmosaic_video_fusion(opt,netG,netM): # clean mosaic print('Clean Mosaic:') length = len(imagepaths) - img_pool = np.zeros((height,width,3*N), dtype='uint8') + + img_pool = [] mosaic_input = np.zeros((INPUT_SIZE,INPUT_SIZE,3*N+1), dtype='uint8') + for i,imagepath in enumerate(imagepaths,0): x,y,size = positions[i][0],positions[i][1],positions[i][2] - + # image read stream mask = cv2.imread(os.path.join(opt.temp_dir+'/mosaic_mask',imagepath),0) if i==0 : for j in range(0,N): - img_pool[:,:,j*3:(j+1)*3] = impro.imread(os.path.join(opt.temp_dir+'/video2image',imagepaths[np.clip(i+j-12,0,len(imagepaths)-1)])) + img_pool.append(impro.imread(os.path.join(opt.temp_dir+'/video2image',imagepaths[np.clip(i+j-12,0,len(imagepaths)-1)]))) else: - img_pool[:,:,0:(N-1)*3] = img_pool[:,:,3:N*3] - img_pool[:,:,(N-1)*3:] = impro.imread(os.path.join(opt.temp_dir+'/video2image',imagepaths[np.clip(i+12,0,len(imagepaths)-1)])) - img_origin = img_pool[:,:,int((N-1)/2)*3:(int((N-1)/2)+1)*3] + img_pool.pop(0) + img_pool.append(impro.imread(os.path.join(opt.temp_dir+'/video2image',imagepaths[np.clip(i+12,0,len(imagepaths)-1)]))) + img_origin = img_pool[12] img_result = img_origin.copy() if size>100: try:#Avoid unknown errors #reshape to network input shape - - mosaic_input[:,:,0:N*3] = impro.resize(img_pool[y-size:y+size,x-size:x+size,:], INPUT_SIZE) + for k in range(N): + mosaic_input[:,:,k*3:(k+1)*3] = impro.resize(img_pool[k][y-size:y+size,x-size:x+size], INPUT_SIZE) mask_input = impro.resize(mask,np.min(img_origin.shape[:2]))[y-size:y+size,x-size:x+size] mosaic_input[:,:,-1] = impro.resize(mask_input, INPUT_SIZE) - mosaic_input_tensor = data.im2tensor(mosaic_input,bgr2rgb=False,use_gpu=opt.use_gpu,use_transform = False,is0_1 = False) - unmosaic_pred = netG(mosaic_input_tensor) + unmosaic_pred = netG(mosaic_input_tensor) img_fake = data.tensor2im(unmosaic_pred,rgb2bgr = False ,is0_1 = False) - img_result = impro.replace_mosaic(img_origin,img_fake,mask,x,y,size,opt.no_feather) + img_result = impro.replace_mosaic(img_origin,img_fake,mask,x,y,size,opt.no_feather) except Exception as e: print('Warning:',e) cv2.imwrite(os.path.join(opt.temp_dir+'/replace_mosaic',imagepath),img_result) diff --git a/cores/options.py b/cores/options.py index 83748ebe68fd02de5b9ad8a4dec2ed8fc93efe6d..7e8c16500ba5d2a07713170991bb75a9d35419bd 100644 --- a/cores/options.py +++ b/cores/options.py @@ -1,5 +1,6 @@ import argparse import os +import sys class Options(): @@ -69,7 +70,7 @@ class Options(): if not os.path.exists(self.opt.media_path): print('Error: Bad media path!') input('Please press any key to exit.\n') - exit(0) + sys.exit(0) if self.opt.mode == 'auto': if 'clean' in model_name or self.opt.traditional: @@ -81,7 +82,7 @@ class Options(): else: print('Please input running model!') input('Please press any key to exit.\n') - exit(0) + sys.exit(0) if self.opt.output_size == 0 and self.opt.mode == 'style': self.opt.output_size = 512 @@ -101,7 +102,7 @@ class Options(): else: print('Type of Generator error!') input('Please press any key to exit.\n') - exit(0) + sys.exit(0) if self.opt.ex_mult == 'auto': if 'face' in model_name: diff --git a/deepmosaic.py b/deepmosaic.py index 132673f2aa763dee6a35e7bd2ede986dc4324c97..7f2b4db4a9d1ab486fd60281696062d6793b154c 100644 --- a/deepmosaic.py +++ b/deepmosaic.py @@ -6,7 +6,8 @@ from util import util from models import loadmodel opt = Options().getparse(test_flag = True) -util.file_init(opt) +if not os.path.isdir(opt.temp_dir): + util.file_init(opt) def main(): @@ -22,8 +23,10 @@ def main(): core.addmosaic_img(opt,netS) elif util.is_video(file): core.addmosaic_video(opt,netS) + util.clean_tempfiles(opt, tmp_init = False) else: print('This type of file is not supported') + util.clean_tempfiles(opt, tmp_init = False) elif opt.mode == 'clean': netM = loadmodel.bisenet(opt,'mosaic') @@ -43,6 +46,7 @@ def main(): core.cleanmosaic_video_fusion(opt,netG,netM) else: core.cleanmosaic_video_byframe(opt,netG,netM) + util.clean_tempfiles(opt, tmp_init = False) else: print('This type of file is not supported') @@ -54,6 +58,7 @@ def main(): core.styletransfer_img(opt,netG) elif util.is_video(file): core.styletransfer_video(opt,netG) + util.clean_tempfiles(opt, tmp_init = False) else: print('This type of file is not supported') @@ -83,5 +88,4 @@ if __name__ == '__main__': print(stack) input('Please press any key to exit.\n') #util.clean_tempfiles(tmp_init = False) - exit(0) - + sys.exit(0) \ No newline at end of file diff --git a/util/util.py b/util/util.py index 565e190631de4cc45a5d0d9da212beafb16171c6..405fd489e8a4f59ac01296d49cfb83bcbe5bedc1 100755 --- a/util/util.py +++ b/util/util.py @@ -70,11 +70,9 @@ def clean_tempfiles(opt,tmp_init=True): os.makedirs(tmpdir) os.makedirs(os.path.join(tmpdir, 'video2image')) os.makedirs(os.path.join(tmpdir, 'addmosaic_image')) - os.makedirs(os.path.join(tmpdir, 'mosaic_crop')) os.makedirs(os.path.join(tmpdir, 'replace_mosaic')) os.makedirs(os.path.join(tmpdir, 'mosaic_mask')) os.makedirs(os.path.join(tmpdir, 'ROI_mask')) - os.makedirs(os.path.join(tmpdir, 'ROI_mask_check')) os.makedirs(os.path.join(tmpdir, 'style_transfer')) def file_init(opt): @@ -90,11 +88,16 @@ def second2stamp(s): s = int(s%60) return "%02d:%02d:%02d" % (h, m, s) -def counttime(t1,t2,now_num,all_num): +def stamp2second(stamp): + substamps = stamp.split(':') + return int(substamps[0])*3600 + int(substamps[1])*60 + int(substamps[2]) + + +def counttime(start_time,current_time,now_num,all_num): ''' - t1,t2: time.time() + start_time,current_time: time.time() ''' - used_time = int(t2-t1) + used_time = int(current_time-start_time) all_time = int(used_time/now_num*all_num) return second2stamp(used_time)+'/'+second2stamp(all_time)