diff --git a/cores/core.py b/cores/core.py index c95d14766c2ff87e112d85fa746561fba85485f0..db992177dd07d3c1a33ea9f924182c6d8c98f3e5 100644 --- a/cores/core.py +++ b/cores/core.py @@ -16,21 +16,20 @@ def video_init(opt,path): fps = opt.fps continue_flag = False + imagepaths = [] + 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: + if imagepaths != []: + imagepaths.sort() + last_frame = int(imagepaths[-1][7:13]) + if (opt.last_time != '00:00:00' and last_frame > fps*(util.stamp2second(opt.last_time)-1)) \ + or (opt.last_time == '00:00:00' and last_frame > 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 not continue_flag: 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) @@ -89,6 +88,7 @@ def addmosaic_video(opt,netS): except Exception as e: print('Warning:',e) cv2.imwrite(os.path.join(opt.temp_dir+'/addmosaic_image',imagepath),img) + os.remove(os.path.join(opt.temp_dir+'/video2image',imagepath)) #preview result and print if not opt.no_preview: @@ -129,6 +129,7 @@ def styletransfer_video(opt,netG): img = impro.imread(os.path.join(opt.temp_dir+'/video2image',imagepath)) img = runmodel.run_styletransfer(opt, netG, img) cv2.imwrite(os.path.join(opt.temp_dir+'/style_transfer',imagepath),img) + os.remove(os.path.join(opt.temp_dir+'/video2image',imagepath)) #preview result and print if not opt.no_preview: @@ -225,6 +226,7 @@ def cleanmosaic_video_byframe(opt,netG,netM): except Exception as e: print('Warning:',e) cv2.imwrite(os.path.join(opt.temp_dir+'/replace_mosaic',imagepath),img_result) + os.remove(os.path.join(opt.temp_dir+'/video2image',imagepath)) #preview result and print if not opt.no_preview: @@ -287,7 +289,8 @@ def cleanmosaic_video_fusion(opt,netG,netM): 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) + cv2.imwrite(os.path.join(opt.temp_dir+'/replace_mosaic',imagepath),img_result) + os.remove(os.path.join(opt.temp_dir+'/video2image',imagepath)) #preview result and print if not opt.no_preview: diff --git a/util/util.py b/util/util.py index 405fd489e8a4f59ac01296d49cfb83bcbe5bedc1..571974e976533b26662d3fb7af1153a24d71b356 100755 --- a/util/util.py +++ b/util/util.py @@ -74,6 +74,9 @@ def clean_tempfiles(opt,tmp_init=True): os.makedirs(os.path.join(tmpdir, 'mosaic_mask')) os.makedirs(os.path.join(tmpdir, 'ROI_mask')) os.makedirs(os.path.join(tmpdir, 'style_transfer')) + # make dataset + os.makedirs(os.path.join(tmpdir, 'mosaic_crop')) + os.makedirs(os.path.join(tmpdir, 'ROI_mask_check')) def file_init(opt): if not os.path.isdir(opt.result_dir):