Temp

cores/core.py

@@ -37,18 +37,30 @@ def addmosaic_video(opt,netS):
     fps,imagepaths = video_init(opt,path)[:2]
     # get position
     positions = []
+    t1 = time.time()
+    if not opt.no_preview:
+        cv2.namedWindow('preview', cv2.WINDOW_NORMAL)
+    
+    print('Find ROI location:')
     for i,imagepath in enumerate(imagepaths,1):
         img = impro.imread(os.path.join('./tmp/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)     
-        print('Find ROI location:')
-        print('\r',str(i)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=35),end='')
+        
+        #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('\nOptimize ROI locations...')
     mask_index = filt.position_medfilt(np.array(positions), 7)
 
     # 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]))
@@ -58,8 +70,17 @@ def addmosaic_video(opt,netS):
             except Exception as e:
                    print('Warning:',e)
         cv2.imwrite(os.path.join('./tmp/addmosaic_image',imagepaths[i]),img)
-        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('preview',mask)
+            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()
+    if not opt.no_preview:
+        cv2.destroyAllWindows()
     ffmpeg.image2video( fps,
                         './tmp/addmosaic_image/output_%06d.'+opt.tempimage_type,
                         './tmp/voice_tmp.mp3',
@@ -80,12 +101,26 @@ def styletransfer_video(opt,netG):
     positions = []
     fps,imagepaths = video_init(opt,path)[:2]
     print('Transfer:')
+    t1 = time.time()
+    if not opt.no_preview:
+        cv2.namedWindow('preview', cv2.WINDOW_NORMAL)
+    length = len(imagepaths)
+
     for i,imagepath in enumerate(imagepaths,1):
         img = impro.imread(os.path.join('./tmp/video2image',imagepath))
         img = runmodel.run_styletransfer(opt, netG, img)
         cv2.imwrite(os.path.join('./tmp/style_transfer',imagepath),img)
-        print('\r',str(i)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=35),end='')
+        
+        #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(length),util.get_bar(100*i/length,num=35),util.counttime(t1,t2,i,len(imagepaths)),end='')
+    
     print()
+    if not opt.no_preview:
+        cv2.destroyAllWindows()
     suffix = os.path.basename(opt.model_path).replace('.pth','').replace('style_','')
     ffmpeg.image2video( fps,
                 './tmp/style_transfer/output_%06d.'+opt.tempimage_type,
@@ -101,24 +136,28 @@ def get_mosaic_positions(opt,netM,imagepaths,savemask=True):
     t1 = time.time()
     if not opt.no_preview:
         cv2.namedWindow('mosaic mask', cv2.WINDOW_NORMAL)
+
     print('Find mosaic location:')
     for i,imagepath in enumerate(imagepaths,1):
         img_origin = impro.imread(os.path.join('./tmp/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)
+        
+        #preview result and print
         if not opt.no_preview:
             cv2.imshow('mosaic mask',mask)
             cv2.waitKey(1) & 0xFF
-        if savemask:
-            cv2.imwrite(os.path.join('./tmp/mosaic_mask',imagepath), mask)
-        positions.append([x,y,size])
         t2 = time.time()
         print('\r',str(i)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=35),util.counttime(t1,t2,i,len(imagepaths)),end='')
+    
     if not opt.no_preview:
         cv2.destroyAllWindows()
     print('\nOptimize mosaic locations...')
     positions =np.array(positions)
     for i in range(3):positions[:,i] = filt.medfilt(positions[:,i],opt.medfilt_num)
-    np.save('./positions.npy', positions)
+
     return positions
 
 def cleanmosaic_img(opt,netG,netM):
@@ -150,6 +189,7 @@ def cleanmosaic_video_byframe(opt,netG,netM):
 
     # clean mosaic
     print('Clean Mosaic:')
+    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))
@@ -166,16 +206,17 @@ def cleanmosaic_video_byframe(opt,netG,netM):
             except Exception as e:
                 print('Warning:',e)
         cv2.imwrite(os.path.join('./tmp/replace_mosaic',imagepath),img_result)
-        #preview result
+        
+        #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(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=35),util.counttime(t1,t2,i+1,len(imagepaths)),end='')
+        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()
-    # to video
     ffmpeg.image2video( fps,
                 './tmp/replace_mosaic/output_%06d.'+opt.tempimage_type,
                 './tmp/voice_tmp.mp3',
@@ -190,10 +231,14 @@ def cleanmosaic_video_fusion(opt,netG,netM):
         INPUT_SIZE = 128
     fps,imagepaths,height,width = video_init(opt,path)
     positions = get_mosaic_positions(opt,netM,imagepaths,savemask=True)
+    t1 = time.time()
+    if not opt.no_preview:
+        cv2.namedWindow('clean', cv2.WINDOW_NORMAL)
     
     # clean mosaic
     print('Clean Mosaic:')
     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]
         
@@ -211,18 +256,19 @@ def cleanmosaic_video_fusion(opt,netG,netM):
         if size>100:
             try:#Avoid unknown errors
                 #reshape to network input shape
-                mosaic_input = np.zeros((INPUT_SIZE,INPUT_SIZE,3*N+1), dtype='uint8')
+                
                 mosaic_input[:,:,0:N*3] = impro.resize(img_pool[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 = data.im2tensor(mosaic_input,bgr2rgb=False,use_gpu=opt.use_gpu,use_transform = False,is0_1 = False)
-                unmosaic_pred = netG(mosaic_input)
+                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)
                 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)        
             except Exception as e:
                 print('Warning:',e)
         cv2.imwrite(os.path.join('./tmp/replace_mosaic',imagepath),img_result)           
+        
         print('\r',str(i+1)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=35),end='')
     print()
     ffmpeg.image2video( fps,