1. Allow unfinished tasks to be restored (#38 #33)

2. Optimize image_pool when processing video

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)           

cores/options.py

 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:

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

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)