1.Add GAN_Feat_loss and VGG_loss when training video HD model

2.Real-time display preview when processing video
3.Input timestamp to determine where to start when processing video
Merge from an anonymous email:
1.Increase digits in frame number to accomodate very long 60fps movies
2.Increase mp3 quality to max (128kbps hurts clarity; 320kbps sounds perfect)
3.Add CLI option to specify a temporary directory
4.Temporary directory should always be a subdirectory named "DeepMosaics" for safety when deleting tree

cores/core.py

@@ -11,13 +11,13 @@ from util import image_processing as impro
 ---------------------Video Init---------------------
 '''
 def video_init(opt,path):
-    util.clean_tempfiles()
+    util.clean_tempfiles(opt)
     fps,endtime,height,width = ffmpeg.get_video_infos(path)
     if opt.fps !=0:
         fps = opt.fps
-    ffmpeg.video2voice(path,'./tmp/voice_tmp.mp3')
-    ffmpeg.video2image(path,'./tmp/video2image/output_%06d.'+opt.tempimage_type,fps)
-    imagepaths=os.listdir('./tmp/video2image')
+    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
 
@@ -35,6 +35,7 @@ def addmosaic_img(opt,netS):
 def addmosaic_video(opt,netS):
     path = opt.media_path
     fps,imagepaths = video_init(opt,path)[:2]
+    length = len(imagepaths)
     # get position
     positions = []
     t1 = time.time()
@@ -43,17 +44,17 @@ def addmosaic_video(opt,netS):
     
     print('Find ROI location:')
     for i,imagepath in enumerate(imagepaths,1):
-        img = impro.imread(os.path.join('./tmp/video2image',imagepath))
+        img = impro.imread(os.path.join(opt.temp_dir+'/video2image',imagepath))
         mask,x,y,size,area = runmodel.get_ROI_position(img,netS,opt)
         positions.append([x,y,area])      
-        cv2.imwrite(os.path.join('./tmp/ROI_mask',imagepath),mask)     
+        cv2.imwrite(os.path.join(opt.temp_dir+'/ROI_mask',imagepath),mask)     
         
         #preview result and print
         if not opt.no_preview:
             cv2.imshow('preview',mask)
             cv2.waitKey(1) & 0xFF
         t2 = time.time()
-        print('\r',str(i)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),util.counttime(t1,t2,i,len(imagepaths)),num=35),end='')
+        print('\r',str(i)+'/'+str(length),util.get_bar(100*i/length,num=35),util.counttime(t1,t2,i,length),end='')
     
     print('\nOptimize ROI locations...')
     mask_index = filt.position_medfilt(np.array(positions), 7)
@@ -61,29 +62,29 @@ def addmosaic_video(opt,netS):
     # add mosaic
     print('Add Mosaic:')
     t1 = time.time()
-    for i in range(len(imagepaths)):
-        mask = impro.imread(os.path.join('./tmp/ROI_mask',imagepaths[mask_index[i]]),'gray')
-        img = impro.imread(os.path.join('./tmp/video2image',imagepaths[i]))
+    for i,imagepath in enumerate(imagepaths,1):
+        mask = impro.imread(os.path.join(opt.temp_dir+'/ROI_mask',imagepaths[mask_index[i-1]]),'gray')
+        img = impro.imread(os.path.join(opt.temp_dir+'/video2image',imagepath))
         if impro.mask_area(mask)>100: 
             try:#Avoid unknown errors
                 img = mosaic.addmosaic(img, mask, opt)
             except Exception as e:
                    print('Warning:',e)
-        cv2.imwrite(os.path.join('./tmp/addmosaic_image',imagepaths[i]),img)
+        cv2.imwrite(os.path.join(opt.temp_dir+'/addmosaic_image',imagepath),img)
         
         #preview result and print
         if not opt.no_preview:
-            cv2.imshow('preview',mask)
+            cv2.imshow('preview',img)
             cv2.waitKey(1) & 0xFF
         t2 = time.time()
-        print('\r',str(i+1)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=35),util.counttime(t1,t2,i,len(imagepaths)),end='')
+        print('\r',str(i)+'/'+str(length),util.get_bar(100*i/length,num=35),util.counttime(t1,t2,i,length),end='')
     
     print()
     if not opt.no_preview:
         cv2.destroyAllWindows()
     ffmpeg.image2video( fps,
-                        './tmp/addmosaic_image/output_%06d.'+opt.tempimage_type,
-                        './tmp/voice_tmp.mp3',
+                        opt.temp_dir+'/addmosaic_image/output_%06d.'+opt.tempimage_type,
+                        opt.temp_dir+'/voice_tmp.mp3',
                          os.path.join(opt.result_dir,os.path.splitext(os.path.basename(path))[0]+'_add.mp4'))
 
 '''
@@ -107,13 +108,13 @@ def styletransfer_video(opt,netG):
     length = len(imagepaths)
 
     for i,imagepath in enumerate(imagepaths,1):
-        img = impro.imread(os.path.join('./tmp/video2image',imagepath))
+        img = impro.imread(os.path.join(opt.temp_dir+'/video2image',imagepath))
         img = runmodel.run_styletransfer(opt, netG, img)
-        cv2.imwrite(os.path.join('./tmp/style_transfer',imagepath),img)
+        cv2.imwrite(os.path.join(opt.temp_dir+'/style_transfer',imagepath),img)
         
         #preview result and print
         if not opt.no_preview:
-            cv2.imshow('preview',mask)
+            cv2.imshow('preview',img)
             cv2.waitKey(1) & 0xFF
         t2 = time.time()
         print('\r',str(i)+'/'+str(length),util.get_bar(100*i/length,num=35),util.counttime(t1,t2,i,len(imagepaths)),end='')
@@ -123,8 +124,8 @@ def styletransfer_video(opt,netG):
         cv2.destroyAllWindows()
     suffix = os.path.basename(opt.model_path).replace('.pth','').replace('style_','')
     ffmpeg.image2video( fps,
-                './tmp/style_transfer/output_%06d.'+opt.tempimage_type,
-                './tmp/voice_tmp.mp3',
+                opt.temp_dir+'/style_transfer/output_%06d.'+opt.tempimage_type,
+                opt.temp_dir+'/voice_tmp.mp3',
                  os.path.join(opt.result_dir,os.path.splitext(os.path.basename(path))[0]+'_'+suffix+'.mp4'))  
 
 '''
@@ -139,11 +140,11 @@ def get_mosaic_positions(opt,netM,imagepaths,savemask=True):
 
     print('Find mosaic location:')
     for i,imagepath in enumerate(imagepaths,1):
-        img_origin = impro.imread(os.path.join('./tmp/video2image',imagepath))
+        img_origin = impro.imread(os.path.join(opt.temp_dir+'/video2image',imagepath))
         x,y,size,mask = runmodel.get_mosaic_position(img_origin,netM,opt)
         positions.append([x,y,size])
         if savemask:
-            cv2.imwrite(os.path.join('./tmp/mosaic_mask',imagepath), mask)
+            cv2.imwrite(os.path.join(opt.temp_dir+'/mosaic_mask',imagepath), mask)
         
         #preview result and print
         if not opt.no_preview:
@@ -192,7 +193,7 @@ def cleanmosaic_video_byframe(opt,netG,netM):
     length = len(imagepaths)
     for i,imagepath in enumerate(imagepaths,0):
         x,y,size = positions[i][0],positions[i][1],positions[i][2]
-        img_origin = impro.imread(os.path.join('./tmp/video2image',imagepath))
+        img_origin = impro.imread(os.path.join(opt.temp_dir+'/video2image',imagepath))
         img_result = img_origin.copy()
         if size > 100:
             try:#Avoid unknown errors
@@ -201,11 +202,11 @@ def cleanmosaic_video_byframe(opt,netG,netM):
                     img_fake = runmodel.traditional_cleaner(img_mosaic,opt)
                 else:
                     img_fake = runmodel.run_pix2pix(img_mosaic,netG,opt)
-                mask = cv2.imread(os.path.join('./tmp/mosaic_mask',imagepath),0)
+                mask = cv2.imread(os.path.join(opt.temp_dir+'/mosaic_mask',imagepath),0)
                 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('./tmp/replace_mosaic',imagepath),img_result)
+        cv2.imwrite(os.path.join(opt.temp_dir+'/replace_mosaic',imagepath),img_result)
         
         #preview result and print
         if not opt.no_preview:
@@ -213,13 +214,12 @@ def cleanmosaic_video_byframe(opt,netG,netM):
             cv2.waitKey(1) & 0xFF
         t2 = time.time()
         print('\r',str(i+1)+'/'+str(length),util.get_bar(100*i/length,num=35),util.counttime(t1,t2,i+1,len(imagepaths)),end='')
-    
     print()
     if not opt.no_preview:
         cv2.destroyAllWindows()
     ffmpeg.image2video( fps,
-                './tmp/replace_mosaic/output_%06d.'+opt.tempimage_type,
-                './tmp/voice_tmp.mp3',
+                opt.temp_dir+'/replace_mosaic/output_%06d.'+opt.tempimage_type,
+                opt.temp_dir+'/voice_tmp.mp3',
                  os.path.join(opt.result_dir,os.path.splitext(os.path.basename(path))[0]+'_clean.mp4'))  
 
 def cleanmosaic_video_fusion(opt,netG,netM):
@@ -237,19 +237,20 @@ 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')
     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('./tmp/mosaic_mask',imagepath),0)
+        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('./tmp/video2image',imagepaths[np.clip(i+j-12,0,len(imagepaths)-1)]))
+                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)]))
         else:
             img_pool[:,:,0:(N-1)*3] = img_pool[:,:,3:N*3]
-            img_pool[:,:,(N-1)*3:] = impro.imread(os.path.join('./tmp/video2image',imagepaths[np.clip(i+12,0,len(imagepaths)-1)]))
+            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_result = img_origin.copy()
 
@@ -267,11 +268,18 @@ 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('./tmp/replace_mosaic',imagepath),img_result)           
+        cv2.imwrite(os.path.join(opt.temp_dir+'/replace_mosaic',imagepath),img_result)           
         
-        print('\r',str(i+1)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=35),end='')
+        #preview result and print
+        if not opt.no_preview:
+            cv2.imshow('clean',img_result)
+            cv2.waitKey(1) & 0xFF
+        t2 = time.time()
+        print('\r',str(i+1)+'/'+str(length),util.get_bar(100*i/length,num=35),util.counttime(t1,t2,i+1,len(imagepaths)),end='')
     print()
+    if not opt.no_preview:
+        cv2.destroyAllWindows()
     ffmpeg.image2video( fps,
-                './tmp/replace_mosaic/output_%06d.'+opt.tempimage_type,
-                './tmp/voice_tmp.mp3',
+                opt.temp_dir+'/replace_mosaic/output_%06d.'+opt.tempimage_type,
+                opt.temp_dir+'/voice_tmp.mp3',
                  os.path.join(opt.result_dir,os.path.splitext(os.path.basename(path))[0]+'_clean.mp4'))        
\ No newline at end of file

cores/options.py

@@ -10,16 +10,19 @@ class Options():
     def initialize(self):
 
         #base
-        self.parser.add_argument('--use_gpu',type=int,default=0, help='if -1, use cpu')
+        self.parser.add_argument('--use_gpu', type=int,default=0, help='if -1, use cpu')
         self.parser.add_argument('--media_path', type=str, default='./imgs/ruoruo.jpg',help='your videos or images path')
+        self.parser.add_argument('-ss', '--start_time', type=str, default='00:00:00',help='start position of video, default is the beginning of video')
+        self.parser.add_argument('-t', '--last_time', type=str, default='00:00:00',help='limit the duration of the video, default is the entire video')
         self.parser.add_argument('--mode', type=str, default='auto',help='Program running mode. auto | add | clean | style')
         self.parser.add_argument('--model_path', type=str, default='./pretrained_models/mosaic/add_face.pth',help='pretrained model path')
         self.parser.add_argument('--result_dir', type=str, default='./result',help='output media will be saved here')
+        self.parser.add_argument('--temp_dir', type=str, default='./tmp', help='Temporary files will go here')
         self.parser.add_argument('--tempimage_type', type=str, default='jpg',help='type of temp image, png | jpg, png is better but occupy more storage space')
         self.parser.add_argument('--netG', type=str, default='auto',
             help='select model to use for netG(Clean mosaic and Transfer style) -> auto | unet_128 | unet_256 | resnet_9blocks | HD | video')
         self.parser.add_argument('--fps', type=int, default=0,help='read and output fps, if 0-> origin')
-        self.parser.add_argument('--no_preview', action='store_true', help='if specified, do not preview images when processing video')
+        self.parser.add_argument('--no_preview', action='store_true', help='if specified,do not preview images when processing video. eg.(when run it on server)')
         self.parser.add_argument('--output_size', type=int, default=0,help='size of output media, if 0 -> origin')
         self.parser.add_argument('--mask_threshold', type=int, default=64,help='threshold of recognize clean or add mosaic position 0~255')
 
@@ -51,8 +54,9 @@ class Options():
         if not self.initialized:
             self.initialize()
         self.opt = self.parser.parse_args()
-
+        
         model_name = os.path.basename(self.opt.model_path)
+        self.opt.temp_dir = os.path.join(self.opt.temp_dir, 'DeepMosaics_temp')
 
         os.environ["CUDA_VISIBLE_DEVICES"] = str(self.opt.use_gpu)
         import torch
@@ -61,6 +65,11 @@ class Options():
         else:
             self.opt.use_gpu = -1
 
+        if not os.path.exists(self.opt.media_path):
+            print('Error: Bad media path!')
+            input('Please press any key to exit.\n')
+            exit(0)
+
         if self.opt.mode == 'auto':
             if 'clean' in model_name or self.opt.traditional:
                 self.opt.mode = 'clean'

deepmosaic.py

@@ -54,15 +54,15 @@ def main():
                 core.styletransfer_img(opt,netG)
             elif util.is_video(file):
                 core.styletransfer_video(opt,netG)
-
             else:
                 print('This type of file is not supported')
 
-    util.clean_tempfiles(tmp_init = False)
+    util.clean_tempfiles(opt, tmp_init = False)
         
 if __name__ == '__main__':
     try:
         main()
+        print('Finished!')
     except Exception as ex:
         print('--------------------ERROR--------------------')
         ex_type, ex_val, ex_stack = sys.exc_info()

docs/options_introduction.md

@@ -7,10 +7,14 @@ If you need more effects,  use '--option your-parameters' to enter what you need
 | :----------: | :------------------------: | :-------------------------------------: |
 |  --use_gpu   |   if -1, do not use gpu    |                    0                    |
 | --media_path | your videos or images path |            ./imgs/ruoruo.jpg            |
+| --start_time | start position of video, default is the beginning of video | '00:00:00' |
+| --last_time | limit the duration of the video, default is the entire video | '00:00:00' |
 |    --mode    |    program running mode(auto/clean/add/style)    |                 'auto'                  |
 | --model_path |   pretrained model path    | ./pretrained_models/mosaic/add_face.pth |
 | --result_dir |  output media will be saved here|                 ./result          |
+| --temp_dir | Temporary files will go here | ./tmp |
 |    --fps    |    read and output fps, if 0-> origin    |                 0                  |
+| --no_preview | if specified,do not preview images when processing video. eg.(when run it on server) | Flase |
 
 ### AddMosaic
 

docs/options_introduction_CN.md

@@ -7,10 +7,14 @@
 | :----------: | :------------------------: | :-------------------------------------: |
 |  --use_gpu   |   if -1, do not use gpu    |                    0                    |
 | --media_path | 需要处理的视频或者照片的路径 |            ./imgs/ruoruo.jpg            |
+| --start_time | 视频开始处理的位置，默认从头开始 | '00:00:00' |
+| --last_time | 处理的视频时长，默认是整个视频 | '00:00:00' |
 |    --mode    |    运行模式(auto/clean/add/style)    |                 'auto'                  |
 | --model_path |   预训练模型的路径    | ./pretrained_models/mosaic/add_face.pth |
 | --result_dir | 保存路径 |                 ./result          |
+| --temp_dir | 临时文件存储目录 | ./tmp |
 |    --fps    |    限制视频输出的fps，0则为默认    |                 0                  |
+| --no_preview | 如果输入，将不会在处理视频时播放实时预览．比如当你在服务器运行的时候 | Flase |
 
 ### 添加马赛克
 

make_datasets/make_video_dataset.py

@@ -52,10 +52,10 @@ for videopath in videopaths:
             timestamps=[]
             fps,endtime,height,width = ffmpeg.get_video_infos(videopath)
             for cut_point in range(1,int((endtime-opt.time)/opt.interval)):
-                util.clean_tempfiles()
-                ffmpeg.video2image(videopath, './tmp/video2image/%05d.'+opt.tempimage_type,fps=1,
+                util.clean_tempfiles(opt)
+                ffmpeg.video2image(videopath, opt.temp_dir+'/video2image/%05d.'+opt.tempimage_type,fps=1,
                     start_time = util.second2stamp(cut_point*opt.interval),last_time = util.second2stamp(opt.time))
-                imagepaths = util.Traversal('./tmp/video2image')
+                imagepaths = util.Traversal(opt.temp_dir+'/video2image')
                 cnt = 0 
                 for i in range(opt.time):
                     img = impro.imread(imagepaths[i])
@@ -80,8 +80,8 @@ for videopath in videopaths:
             util.makedirs(origindir)
             util.makedirs(maskdir)
 
-            util.clean_tempfiles()
-            ffmpeg.video2image(videopath, './tmp/video2image/%05d.'+opt.tempimage_type,
+            util.clean_tempfiles(opt)
+            ffmpeg.video2image(videopath, opt.temp_dir+'/video2image/%05d.'+opt.tempimage_type,
                 start_time = timestamp,last_time = util.second2stamp(opt.time))
             
             endtime = datetime.datetime.now()
@@ -89,7 +89,7 @@ for videopath in videopaths:
                 util.get_bar(100*video_cnt/len(videopaths),35),'',
                 util.second2stamp((endtime-starttime).seconds)+'/'+util.second2stamp((endtime-starttime).seconds/video_cnt*len(videopaths)))
 
-            imagepaths = util.Traversal('./tmp/video2image')
+            imagepaths = util.Traversal(opt.temp_dir+'/video2image')
             imagepaths = sorted(imagepaths)
             imgs=[];masks=[]
             mask_flag = False

util/ffmpeg.py

@@ -2,35 +2,60 @@ import os,json
 
 # ffmpeg 3.4.6
 
-def video2image(videopath,imagepath,fps=0,start_time=0,last_time=0):
-    if start_time == 0:
-        if fps == 0:
-            os.system('ffmpeg -i "'+videopath+'" -f image2 '+'-q:v -0 '+imagepath)
-        else:
-            os.system('ffmpeg -i "'+videopath+'" -r '+str(fps)+' -f image2 '+'-q:v -0 '+imagepath)
-    else:
-        if fps == 0:
-            os.system('ffmpeg -ss '+start_time+' -t '+last_time+' -i "'+videopath+'" -f image2 '+'-q:v -0 '+imagepath)
-        else:
-            os.system('ffmpeg -ss '+start_time+' -t '+last_time+' -i "'+videopath+'" -r '+str(fps)+' -f image2 '+'-q:v -0 '+imagepath)
-
-
-def video2voice(videopath,voicepath):
-    os.system('ffmpeg -i "'+videopath+'" -f mp3 '+voicepath)
+def args2cmd(args):
+    cmd = ''
+    for arg in args:
+        cmd += (arg+' ')
+    return cmd
+
+def run(args,mode = 0):
+
+    if mode == 0:
+        cmd = args2cmd(args)
+        os.system(cmd)
+
+    elif mode == 1:
+        '''
+        out_string = os.popen(cmd_str).read()
+        For chinese path in Windows
+        https://blog.csdn.net/weixin_43903378/article/details/91979025
+        '''
+        cmd = args2cmd(args)
+        stream = os.popen(cmd)._stream
+        sout = stream.buffer.read().decode(encoding='utf-8')
+        return sout
+
+    elif mode == 2:
+        cmd = args2cmd(args)
+        p = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        sout = p.stdout.readlines()
+        return sout
+
+def video2image(videopath, imagepath, fps=0, start_time='00:00:00', last_time='00:00:00'):
+    args = ['ffmpeg', '-i', '"'+videopath+'"']
+    if last_time != '00:00:00':
+        args += ['-ss', start_time]
+        args += ['-t', last_time]
+    if fps != 0:
+        args += ['-r', str(fps)]
+    args += ['-f', 'image2','-q:v','-0',imagepath]
+    run(args)
+
+def video2voice(videopath, voicepath, start_time='00:00:00', last_time='00:00:00'):
+    args = ['ffmpeg', '-i', '"'+videopath+'"','-f mp3','-b:a 320k']
+    if last_time != '00:00:00':
+        args += ['-ss', start_time]
+        args += ['-t', last_time]
+    args += [voicepath]
+    run(args)
 
 def image2video(fps,imagepath,voicepath,videopath):
-    os.system('ffmpeg -y -r '+str(fps)+' -i '+imagepath+' -vcodec libx264 '+'./tmp/video_tmp.mp4')
-    #os.system('ffmpeg -f image2 -i '+imagepath+' -vcodec libx264 -r '+str(fps)+' ./tmp/video_tmp.mp4')
-    os.system('ffmpeg -i ./tmp/video_tmp.mp4 -i "'+voicepath+'" -vcodec copy -acodec copy '+videopath)
+    os.system('ffmpeg -y -r '+str(fps)+' -i '+imagepath+' -vcodec libx264 '+os.path.split(voicepath)[0]+'/video_tmp.mp4')
+    os.system('ffmpeg -i '+os.path.split(voicepath)[0]+'/video_tmp.mp4'+' -i "'+voicepath+'" -vcodec copy -acodec aac '+videopath)
 
 def get_video_infos(videopath):
-    cmd_str =  'ffprobe -v quiet -print_format json -show_format -show_streams -i "' + videopath + '"'  
-    #out_string = os.popen(cmd_str).read()
-    #For chinese path in Windows
-    #https://blog.csdn.net/weixin_43903378/article/details/91979025
-    stream = os.popen(cmd_str)._stream
-    out_string = stream.buffer.read().decode(encoding='utf-8')
-
+    args =  ['ffprobe -v quiet -print_format json -show_format -show_streams', '-i', '"'+videopath+'"']
+    out_string = run(args,mode=1)
     infos = json.loads(out_string)
     try:
         fps = eval(infos['streams'][0]['avg_frame_rate'])

util/util.py

@@ -61,26 +61,27 @@ def makedirs(path):
         os.makedirs(path)
         print('makedir:',path)
 
-def clean_tempfiles(tmp_init=True):
-    if os.path.isdir('./tmp'): 
+def clean_tempfiles(opt,tmp_init=True):
+    tmpdir = opt.temp_dir
+    if os.path.isdir(tmpdir): 
         print('Clean temp...')  
-        shutil.rmtree('./tmp')
+        shutil.rmtree(tmpdir)
     if tmp_init:
-        os.makedirs('./tmp')
-        os.makedirs('./tmp/video2image')
-        os.makedirs('./tmp/addmosaic_image')
-        os.makedirs('./tmp/mosaic_crop')
-        os.makedirs('./tmp/replace_mosaic')
-        os.makedirs('./tmp/mosaic_mask')
-        os.makedirs('./tmp/ROI_mask')
-        os.makedirs('./tmp/ROI_mask_check')
-        os.makedirs('./tmp/style_transfer')
-
+        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):
     if not os.path.isdir(opt.result_dir):
         os.makedirs(opt.result_dir)
         print('makedir:',opt.result_dir)
-    clean_tempfiles(True)
+    clean_tempfiles(opt,True)
 
 def second2stamp(s):
     h = int(s/3600)