From e31088f5977054f41ff73b1f474bd04c6c3a9eed Mon Sep 17 00:00:00 2001 From: hypox64 Date: Sat, 19 Dec 2020 14:06:01 +0800 Subject: [PATCH] Clean cache during processing --- cores/core.py | 31 +++++++++++++++++-------------- util/util.py | 3 +++ 2 files changed, 20 insertions(+), 14 deletions(-) diff --git a/cores/core.py b/cores/core.py index c95d147..db99217 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 405fd48..571974e 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):