Merge branch 'dev'

.gitignore

@@ -142,8 +142,9 @@ test*
 video_tmp/
 result/
 nohup.out
+.vscode/
+
 #./
-/.vscode
 /pix2pix
 /pix2pixHD
 /tmp
@@ -157,6 +158,7 @@ nohup.out
 /deepmosaic_window
 /sftp-config.json
 /exe
+
 #./make_datasets
 /make_datasets/video
 /make_datasets/tmp
@@ -172,6 +174,7 @@ nohup.out
 #mediafile
 *iter
 *.pth
+*.pt
 *.jpeg
 *.bmp
 *.mp4
@@ -185,4 +188,39 @@ nohup.out
 *.JPEG
 *.exe
 *.npy
-*.psd
\ No newline at end of file
+*.psd
+
+
+##############################cpp###################################
+# Prerequisites
+*.d
+
+# Compiled Object files
+*.slo
+*.lo
+*.o
+*.obj
+
+# Precompiled Headers
+*.gch
+*.pch
+
+# Compiled Dynamic libraries
+*.so
+*.dylib
+*.dll
+
+# Fortran module files
+*.mod
+*.smod
+
+# Compiled Static libraries
+*.lai
+*.la
+*.a
+*.lib
+
+# Executables
+*.exe
+*.out
+*.app
\ No newline at end of file

README.md

@@ -9,6 +9,7 @@ You can use it to automatically remove the mosaics in images and videos, or add
 
 ### Examples
 ![image](./imgs/hand.gif)
+
 origin | auto add mosaic |  auto clean mosaic  
 :-:|:-:|:-:
 ![image](./imgs/example/lena.jpg) | ![image](./imgs/example/lena_add.jpg) | ![image](./imgs/example/lena_clean.jpg) 
@@ -33,7 +34,7 @@ An interesting example:[Ricardo Milos to cat](https://www.bilibili.com/video/BV1
 You can either run DeepMosaics via a pre-built binary package, or from source.<br>
 
 ### Try it on web
-You can simply try to remove the mosaic on the face at this [website](http://118.89.27.46:5000/).<br>
+You can simply try to remove the mosaic on the **face** at this [website](http://118.89.27.46:5000/).<br>
 ### Pre-built binary package
 For Windows, we bulid a GUI version for easy testing.<br>
 Download this version, and a pre-trained model via [[Google Drive]](https://drive.google.com/open?id=1LTERcN33McoiztYEwBxMuRjjgxh4DEPs)  [[百度云,提取码1x0a]](https://pan.baidu.com/s/10rN3U3zd5TmfGpO_PEShqQ) <br>
@@ -61,7 +62,7 @@ Attentions:<br>
 This code depends on opencv-python, torchvision available via pip install.
 #### Clone this repo
 ```bash
-git clone https://github.com/HypoX64/DeepMosaics
+git clone https://github.com/HypoX64/DeepMosaics.git
 cd DeepMosaics
 ```
 #### Get Pre-Trained Models

README_CN.md

@@ -10,6 +10,7 @@
 
 ### 例子
 ![image](./imgs/hand.gif)
+
 原始 | 自动打码 |  自动去码  
 :-:|:-:|:-:
 ![image](./imgs/example/lena.jpg) | ![image](./imgs/example/lena_add.jpg) | ![image](./imgs/example/lena_clean.jpg) 
@@ -33,7 +34,7 @@
 ## 如何运行
 可以通过我们预编译好的二进制包或源代码运行.<br>
 ### 在网页中运行
-打开[这个网站](http://118.89.27.46:5000/)上传照片，将获得去除马赛克后的结果，受限与当地法律，目前只支持人脸.<br>
+打开[这个网站](http://118.89.27.46:5000/)上传照片，将获得去除马赛克后的结果，受限于当地法律，**目前只支持人脸**.<br>
 ### 预编译的程序包
 对于Windows用户,我们提供了包含GUI界面的免安装软件包.<br>
 可以通过下面两种方式进行下载: [[Google Drive]](https://drive.google.com/open?id=1LTERcN33McoiztYEwBxMuRjjgxh4DEPs)  [[百度云,提取码1x0a]](https://pan.baidu.com/s/10rN3U3zd5TmfGpO_PEShqQ) <br>
@@ -57,10 +58,10 @@
   - [Pytorch 1.0+](https://pytorch.org/) 
   - CPU or NVIDIA GPU + CUDA CuDNN<br>
 #### Python依赖项
-代码依赖于opencv-python以及 torchvision,可有通过pip install 进行安装.
+代码依赖于opencv-python以及 torchvision,可以通过pip install 进行安装.
 #### 克隆源代码
 ```bash
-git clone https://github.com/HypoX64/DeepMosaics
+git clone https://github.com/HypoX64/DeepMosaics.git
 cd DeepMosaics
 ```
 #### 下载预训练模型

cores/add.py

+import os
+from queue import Queue
+from threading import Thread
+import time
+import numpy as np
+import cv2
+from models import runmodel
+from util import mosaic,util,ffmpeg,filt
+from util import image_processing as impro
+from .init import video_init
+
+
+'''
+---------------------Add Mosaic---------------------
+'''
+def addmosaic_img(opt,netS):
+    path = opt.media_path
+    print('Add Mosaic:',path)
+    img = impro.imread(path)
+    mask = runmodel.get_ROI_position(img,netS,opt)[0]
+    img = mosaic.addmosaic(img,mask,opt)
+    impro.imwrite(os.path.join(opt.result_dir,os.path.splitext(os.path.basename(path))[0]+'_add.jpg'),img)
+
+def get_roi_positions(opt,netS,imagepaths,savemask=True):
+    # resume
+    continue_flag = False
+    if os.path.isfile(os.path.join(opt.temp_dir,'step.json')):
+        step = util.loadjson(os.path.join(opt.temp_dir,'step.json'))
+        resume_frame = int(step['frame'])
+        if int(step['step'])>2:
+            mask_index = np.load(os.path.join(opt.temp_dir,'mask_index.npy'))
+            return mask_index
+        if int(step['step'])>=2 and resume_frame>0:
+            pre_positions = np.load(os.path.join(opt.temp_dir,'roi_positions.npy'))
+            continue_flag = True
+            imagepaths = imagepaths[resume_frame:]
+            
+    positions = []
+    t1 = time.time()
+    if not opt.no_preview:
+        cv2.namedWindow('mask', cv2.WINDOW_NORMAL)
+    print('Step:2/4 -- Find mosaic location')
+
+    img_read_pool = Queue(4)
+    def loader(imagepaths):
+        for imagepath in imagepaths:
+            img_origin = impro.imread(os.path.join(opt.temp_dir+'/video2image',imagepath))
+            img_read_pool.put(img_origin)
+    t = Thread(target=loader,args=(imagepaths,))
+    t.daemon = True
+    t.start()
+
+    for i,imagepath in enumerate(imagepaths,1):
+        img_origin = img_read_pool.get()
+        mask,x,y,size,area = runmodel.get_ROI_position(img_origin,netS,opt)
+        positions.append([x,y,area])  
+        if savemask:
+            t = Thread(target=cv2.imwrite,args=(os.path.join(opt.temp_dir+'/ROI_mask',imagepath), mask,))
+            t.start()
+        if i%1000==0:
+            save_positions = np.array(positions)
+            if continue_flag:
+                save_positions = np.concatenate((pre_positions,save_positions),axis=0)
+            np.save(os.path.join(opt.temp_dir,'roi_positions.npy'),save_positions)
+            step = {'step':2,'frame':i+resume_frame}
+            util.savejson(os.path.join(opt.temp_dir,'step.json'),step)
+
+        #preview result and print
+        if not opt.no_preview:
+            cv2.imshow('mask',mask)
+            cv2.waitKey(1) & 0xFF
+        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 ROI locations...')
+    if continue_flag:
+        positions = np.concatenate((pre_positions,positions),axis=0)
+    mask_index = filt.position_medfilt(np.array(positions), 7)
+    step = {'step':3,'frame':0}
+    util.savejson(os.path.join(opt.temp_dir,'step.json'),step)
+    np.save(os.path.join(opt.temp_dir,'roi_positions.npy'),positions)
+    np.save(os.path.join(opt.temp_dir,'mask_index.npy'),np.array(mask_index))
+
+    return mask_index
+
+def addmosaic_video(opt,netS):
+    path = opt.media_path
+    fps,imagepaths = video_init(opt,path)[:2]
+    length = len(imagepaths)
+    start_frame = int(imagepaths[0][7:13])
+    mask_index = get_roi_positions(opt,netS,imagepaths)[(start_frame-1):]
+
+    t1 = time.time()
+    if not opt.no_preview:
+        cv2.namedWindow('preview', cv2.WINDOW_NORMAL)
+
+    # add mosaic
+    print('Step:3/4 -- Add Mosaic:')
+    t1 = time.time()
+    # print(mask_index)
+    for i,imagepath in enumerate(imagepaths,1):
+        mask = impro.imread(os.path.join(opt.temp_dir+'/ROI_mask',imagepaths[np.clip(mask_index[i-1]-start_frame,0,1000000)]),'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)
+        t = Thread(target=cv2.imwrite,args=(os.path.join(opt.temp_dir+'/addmosaic_image',imagepath),img))
+        t.start()
+        os.remove(os.path.join(opt.temp_dir+'/video2image',imagepath))
+        
+        #preview result and print
+        if not opt.no_preview:
+            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,length),end='')
+    
+    print()
+    if not opt.no_preview:
+        cv2.destroyAllWindows()
+    print('Step:4/4 -- Convert images to video')
+    ffmpeg.image2video( fps,
+                        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'))
\ No newline at end of file

cores/core.py → cores/clean.py

 import os
 import time
-import torch
 import numpy as np
 import cv2
-
-from models import runmodel,loadmodel
-from util import mosaic,util,ffmpeg,filt,data
+import torch
+from models import runmodel
+from util import data,util,ffmpeg,filt
 from util import image_processing as impro
-
-'''
----------------------Video Init---------------------
-'''
-def video_init(opt,path):
-    fps,endtime,height,width = ffmpeg.get_video_infos(path)
-    if opt.fps !=0:
-        fps = opt.fps
-
-    continue_flag = False
-    imagepaths = []
-
-    if os.path.isdir(opt.temp_dir):
-        imagepaths = os.listdir(opt.temp_dir+'/video2image')
-        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:
-        print('Step:1/4 -- Convert video to images')
-        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
-
-'''
----------------------Add Mosaic---------------------
-'''
-def addmosaic_img(opt,netS):
-    path = opt.media_path
-    print('Add Mosaic:',path)
-    img = impro.imread(path)
-    mask = runmodel.get_ROI_position(img,netS,opt)[0]
-    img = mosaic.addmosaic(img,mask,opt)
-    impro.imwrite(os.path.join(opt.result_dir,os.path.splitext(os.path.basename(path))[0]+'_add.jpg'),img)
-
-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()
-    if not opt.no_preview:
-        cv2.namedWindow('preview', cv2.WINDOW_NORMAL)
-    
-    print('Step:2/4 -- Find ROI location')
-    for i,imagepath in enumerate(imagepaths,1):
-        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(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(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)
-
-    # add mosaic
-    print('Step:3/4 -- Add Mosaic:')
-    t1 = time.time()
-    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(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:
-            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,length),end='')
-    
-    print()
-    if not opt.no_preview:
-        cv2.destroyAllWindows()
-    print('Step:4/4 -- Convert images to video')
-    ffmpeg.image2video( fps,
-                        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'))
-
-'''
----------------------Style Transfer---------------------
-'''
-def styletransfer_img(opt,netG):
-    print('Style Transfer_img:',opt.media_path)
-    img = impro.imread(opt.media_path)
-    img = runmodel.run_styletransfer(opt, netG, img)
-    suffix = os.path.basename(opt.model_path).replace('.pth','').replace('style_','')
-    impro.imwrite(os.path.join(opt.result_dir,os.path.splitext(os.path.basename(opt.media_path))[0]+'_'+suffix+'.jpg'),img)
-
-def styletransfer_video(opt,netG):
-    path = opt.media_path
-    positions = []
-    fps,imagepaths = video_init(opt,path)[:2]
-    print('Step:2/4 -- 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(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:
-            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='')
-    
-    print()
-    if not opt.no_preview:
-        cv2.destroyAllWindows()
-    suffix = os.path.basename(opt.model_path).replace('.pth','').replace('style_','')
-    print('Step:4/4 -- Convert images to video')
-    ffmpeg.image2video( fps,
-                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'))  
+from .init import video_init
+from multiprocessing import Queue, Process
+from threading import Thread
 
 '''
 ---------------------Clean Mosaic---------------------
 '''
 def get_mosaic_positions(opt,netM,imagepaths,savemask=True):
-    # get mosaic position
+    # resume
+    continue_flag = False
+    if os.path.isfile(os.path.join(opt.temp_dir,'step.json')):
+        step = util.loadjson(os.path.join(opt.temp_dir,'step.json'))
+        resume_frame = int(step['frame'])
+        if int(step['step'])>2:
+            pre_positions = np.load(os.path.join(opt.temp_dir,'mosaic_positions.npy'))
+            return pre_positions
+        if int(step['step'])>=2 and resume_frame>0:
+            pre_positions = np.load(os.path.join(opt.temp_dir,'mosaic_positions.npy'))
+            continue_flag = True
+            imagepaths = imagepaths[resume_frame:]
+            
     positions = []
     t1 = time.time()
     if not opt.no_preview:
         cv2.namedWindow('mosaic mask', cv2.WINDOW_NORMAL)
     print('Step:2/4 -- Find mosaic location')
+
+    img_read_pool = Queue(4)
+    def loader(imagepaths):
+        for imagepath in imagepaths:
+            img_origin = impro.imread(os.path.join(opt.temp_dir+'/video2image',imagepath))
+            img_read_pool.put(img_origin)
+    t = Thread(target=loader,args=(imagepaths,))
+    t.setDaemon(True)
+    t.start()
+
     for i,imagepath in enumerate(imagepaths,1):
-        img_origin = impro.imread(os.path.join(opt.temp_dir+'/video2image',imagepath))
+        img_origin = img_read_pool.get()
         x,y,size,mask = runmodel.get_mosaic_position(img_origin,netM,opt)
         positions.append([x,y,size])
         if savemask:
-            cv2.imwrite(os.path.join(opt.temp_dir+'/mosaic_mask',imagepath), mask)
-        
+            t = Thread(target=cv2.imwrite,args=(os.path.join(opt.temp_dir+'/mosaic_mask',imagepath), mask,))
+            t.start()
+        if i%1000==0:
+            save_positions = np.array(positions)
+            if continue_flag:
+                save_positions = np.concatenate((pre_positions,save_positions),axis=0)
+            np.save(os.path.join(opt.temp_dir,'mosaic_positions.npy'),save_positions)
+            step = {'step':2,'frame':i+resume_frame}
+            util.savejson(os.path.join(opt.temp_dir,'step.json'),step)
+
         #preview result and print
         if not opt.no_preview:
             cv2.imshow('mosaic mask',mask)
@@ -179,7 +68,12 @@ def get_mosaic_positions(opt,netM,imagepaths,savemask=True):
         cv2.destroyAllWindows()
     print('\nOptimize mosaic locations...')
     positions =np.array(positions)
+    if continue_flag:
+        positions = np.concatenate((pre_positions,positions),axis=0)
     for i in range(3):positions[:,i] = filt.medfilt(positions[:,i],opt.medfilt_num)
+    step = {'step':3,'frame':0}
+    util.savejson(os.path.join(opt.temp_dir,'step.json'),step)
+    np.save(os.path.join(opt.temp_dir,'mosaic_positions.npy'),positions)
 
     return positions
 
@@ -216,8 +110,10 @@ def cleanmosaic_img_server(opt,img_origin,netG,netM):
 
 def cleanmosaic_video_byframe(opt,netG,netM):
     path = opt.media_path
-    fps,imagepaths = video_init(opt,path)[:2]
-    positions = get_mosaic_positions(opt,netM,imagepaths,savemask=True)
+    fps,imagepaths,height,width = video_init(opt,path)
+    start_frame = int(imagepaths[0][7:13])
+    positions = get_mosaic_positions(opt,netM,imagepaths,savemask=True)[(start_frame-1):]
+
     t1 = time.time()
     if not opt.no_preview:
         cv2.namedWindow('clean', cv2.WINDOW_NORMAL)
@@ -240,7 +136,8 @@ def cleanmosaic_video_byframe(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)
+        t = Thread(target=cv2.imwrite,args=(os.path.join(opt.temp_dir+'/replace_mosaic',imagepath), img_result,))
+        t.start()
         os.remove(os.path.join(opt.temp_dir+'/video2image',imagepath))
         
         #preview result and print
@@ -270,7 +167,8 @@ def cleanmosaic_video_fusion(opt,netG,netM):
     init_flag = True
     
     fps,imagepaths,height,width = video_init(opt,path)
-    positions = get_mosaic_positions(opt,netM,imagepaths,savemask=True)
+    start_frame = int(imagepaths[0][7:13])
+    positions = get_mosaic_positions(opt,netM,imagepaths,savemask=True)[(start_frame-1):]
     t1 = time.time()
     if not opt.no_preview:
         cv2.namedWindow('clean', cv2.WINDOW_NORMAL)
@@ -278,7 +176,24 @@ def cleanmosaic_video_fusion(opt,netG,netM):
     # clean mosaic
     print('Step:3/4 -- Clean Mosaic:')
     length = len(imagepaths)
-    
+    write_pool = Queue(4)
+    show_pool = Queue(4)
+    def write_result():
+        while True:
+            save_ori,imagepath,img_origin,img_fake,x,y,size = write_pool.get()
+            if save_ori:
+                img_result = img_origin
+            else:
+                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)
+            if not opt.no_preview:
+                show_pool.put(img_result.copy())
+            cv2.imwrite(os.path.join(opt.temp_dir+'/replace_mosaic',imagepath),img_result)
+            os.remove(os.path.join(opt.temp_dir+'/video2image',imagepath))
+    t = Thread(target=write_result,args=())
+    t.setDaemon(True)
+    t.start()
+
     for i,imagepath in enumerate(imagepaths,0):
         x,y,size = positions[i][0],positions[i][1],positions[i][2]
         input_stream = []
@@ -290,12 +205,17 @@ def cleanmosaic_video_fusion(opt,netG,netM):
             img_pool.pop(0)
             img_pool.append(impro.imread(os.path.join(opt.temp_dir+'/video2image',imagepaths[np.clip(i+LEFT_FRAME,0,len(imagepaths)-1)])))
         img_origin = img_pool[LEFT_FRAME]
-        img_result = img_origin.copy()
+
+        # preview result and print
+        if not opt.no_preview:
+            if show_pool.qsize()>3:   
+                cv2.imshow('clean',show_pool.get())
+                cv2.waitKey(1) & 0xFF
 
         if size>50:
             try:#Avoid unknown errors
                 for pos in FRAME_POS:
-                    input_stream.append(impro.resize(img_pool[pos][y-size:y+size,x-size:x+size], INPUT_SIZE)[:,:,::-1])
+                    input_stream.append(impro.resize(img_pool[pos][y-size:y+size,x-size:x+size], INPUT_SIZE,interpolation=cv2.INTER_CUBIC)[:,:,::-1])
                 if init_flag:
                     init_flag = False
                     previous_frame = input_stream[N]
@@ -307,28 +227,23 @@ def cleanmosaic_video_fusion(opt,netG,netM):
                     unmosaic_pred = netG(input_stream,previous_frame)
                 img_fake = data.tensor2im(unmosaic_pred,rgb2bgr = True)
                 previous_frame = unmosaic_pred
-                # previous_frame = data.tensor2im(unmosaic_pred,rgb2bgr = True)
-                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)
+                write_pool.put([False,imagepath,img_origin.copy(),img_fake.copy(),x,y,size])
             except Exception as e:
                 init_flag = True
                 print('Error:',e)
         else:
+            write_pool.put([True,imagepath,img_origin.copy(),-1,-1,-1,-1])
             init_flag = True
-        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:
-            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()
+    write_pool.close()
+    show_pool.close()
     if not opt.no_preview:
         cv2.destroyAllWindows()
     print('Step:4/4 -- Convert images to video')
     ffmpeg.image2video( fps,
                 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
+                 os.path.join(opt.result_dir,os.path.splitext(os.path.basename(path))[0]+'_clean.mp4')) 
\ No newline at end of file

cores/init.py

+import os
+from util import util,ffmpeg
+
+'''
+---------------------Video Init---------------------
+'''
+def video_init(opt,path):
+    fps,endtime,height,width = ffmpeg.get_video_infos(path)
+    if opt.fps !=0:
+        fps = opt.fps
+
+    # resume
+    if os.path.isfile(os.path.join(opt.temp_dir,'step.json')):
+        step = util.loadjson(os.path.join(opt.temp_dir,'step.json'))
+        if int(step['step'])>=1:
+            choose = input('There is an unfinished video. Continue it? [y/n] ')
+            if choose.lower() =='yes' or choose.lower() == 'y':
+                imagepaths = os.listdir(opt.temp_dir+'/video2image')
+                imagepaths.sort()
+                return fps,imagepaths,height,width
+    
+    print('Step:1/4 -- Convert video to images')
+    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()
+    step = {'step':2,'frame':0}
+    util.savejson(os.path.join(opt.temp_dir,'step.json'),step)
+
+    return fps,imagepaths,height,width
\ No newline at end of file

cores/options.py

@@ -26,7 +26,7 @@ class Options():
         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. 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')
+        self.parser.add_argument('--mask_threshold', type=int, default=64,help='Mosaic detection threshold (0~255). The smaller is it, the more likely judged as a mosaic area.')
 
         #AddMosaic
         self.parser.add_argument('--mosaic_mod', type=str, default='squa_avg',help='type of mosaic -> squa_avg | squa_random | squa_avg_circle_edge | rect_avg | random')

cores/style.py

+import os
+import time
+import numpy as np
+import cv2
+from models import runmodel
+from util import mosaic,util,ffmpeg,filt
+from util import image_processing as impro
+from .init import video_init
+
+'''
+---------------------Style Transfer---------------------
+'''
+def styletransfer_img(opt,netG):
+    print('Style Transfer_img:',opt.media_path)
+    img = impro.imread(opt.media_path)
+    img = runmodel.run_styletransfer(opt, netG, img)
+    suffix = os.path.basename(opt.model_path).replace('.pth','').replace('style_','')
+    impro.imwrite(os.path.join(opt.result_dir,os.path.splitext(os.path.basename(opt.media_path))[0]+'_'+suffix+'.jpg'),img)
+
+def styletransfer_video(opt,netG):
+    path = opt.media_path
+    fps,imagepaths = video_init(opt,path)[:2]
+    print('Step:2/4 -- 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(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:
+            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='')
+    
+    print()
+    if not opt.no_preview:
+        cv2.destroyAllWindows()
+    suffix = os.path.basename(opt.model_path).replace('.pth','').replace('style_','')
+    print('Step:4/4 -- Convert images to video')
+    ffmpeg.image2video( fps,
+                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')) 
\ No newline at end of file

cpp/CMakeLists.txt

+cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
+set(CMAKE_CXX_STANDARD 14)
+
+project(DeepMosaics)
+set(LIBRARY_OUTPUT_PATH ${PROJECT_SOURCE_DIR}/lib) #链接库路径
+
+set(Torch_DIR /home/hypo/libtorch/share/cmake/Torch)
+find_package(Torch REQUIRED)
+
+set(OpenCV_DIR /home/hypo/opencv-4.4.0)
+find_package(OpenCV REQUIRED)
+
+# Add sub directories
+add_subdirectory(example)
+add_subdirectory(utils)
+
+# set_property(TARGET ${PROJECT_NAME} PROPERTY CXX_STANDARD 14)
+# cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
+# project(main)
+# set(LIBRARY_OUTPUT_PATH ${PROJECT_SOURCE_DIR}/lib) #链接库路径
+
+# set(Torch_DIR /home/hypo/libtorch/share/cmake/Torch)
+# find_package(Torch REQUIRED)
+
+# set(OpenCV_DIR /home/hypo/opencv-4.4.0)
+# find_package(OpenCV REQUIRED)
+
+# # 查找当前目录下的所有源文件
+# # 并将名称保存到 DIR_SRCS 变量
+# # aux_source_directory(. DIR_SRCS)
+# add_subdirectory(utils)
+
+# add_executable(main main.cpp)
+# # target_link_libraries(main )
+# # include_directories( "${OpenCV_INCLUDE_DIRS}" )
+# target_link_libraries( main "${TORCH_LIBRARIES}" "${OpenCV_LIBS}" utils)
+
+# set_property(TARGET main PROPERTY CXX_STANDARD 14)

cpp/README.md

+### C++ version for DeepMosaics
+* I am learning c++ through this project...
+* It is under development...
\ No newline at end of file

cpp/example/CMakeLists.txt

+# project(example)
+# add_executable("${PROJECT_NAME}" deepmosaic.cpp)
+# target_link_libraries(  "${PROJECT_NAME}" 
+#                         "${TORCH_LIBRARIES}" 
+#                         "${OpenCV_LIBS}"
+#                         utils)
+
+file(GLOB_RECURSE srcs RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}" "${CMAKE_CURRENT_SOURCE_DIR}/*.cpp")
+foreach(sourcefile IN LISTS srcs)
+    string( REPLACE ".cpp" "" binname ${sourcefile})
+    add_executable( ${binname} ${sourcefile} )
+    target_link_libraries(  ${binname} 
+                            "${TORCH_LIBRARIES}" 
+                            "${OpenCV_LIBS}"
+                            utils)
+    # set_property(TARGET ${binname} PROPERTY CXX_STANDARD 14)
+endforeach()
\ No newline at end of file

cpp/example/deepmosaic.cpp

+#include <stdio.h>
+#include <string.h>
+#include <unistd.h>
+#include <iostream>
+#include <list>
+#include <vector>
+
+#include <torch/script.h>
+#include <torch/torch.h>
+#include <opencv2/opencv.hpp>
+
+#include "data.hpp"
+#include "util.hpp"
+
+int main() {
+    std::string path = util::current_path();
+
+    std::string net_path = "../res/models/mosaic_position.pth";
+    std::string img_path = "../res/test_media/face/d.jpg";
+
+    cv::Mat img = cv::imread(img_path);
+    cv::resize(img, img, cv::Size(360, 360), 2);
+    // img.convertTo(img, CV_32F);
+    torch::Tensor img_tensor =
+        torch::from_blob(img.data, {1, img.rows, img.cols, 3}, torch::kByte);
+    img_tensor = img_tensor.permute({0, 3, 1, 2});
+    img_tensor = img_tensor.toType(torch::kFloat);
+    img_tensor = img_tensor.div(255);
+    std::cout << img_tensor.sizes() << "\n";
+
+    // end = clock();
+    // dur = (double)(end - start);
+    // printf("Use Time:%f\n", (dur / CLOCKS_PER_SEC));
+
+    // std::string net_path = "../res/models/mosaic_position.pt";
+    // torch::jit::script::Module net;
+    // try{
+    //     // if (!isfile(net_path)){
+    //     //     std::cerr<<"model does not exist\n";
+    //     // }
+
+    //     net = torch::jit::load(net_path);
+    // }
+    // catch(const std::exception& e){
+    //     std::cerr << "error loading the model\n";
+    //     return -1;
+    // }
+
+    // torch::Tensor example = torch::ones({1,3,360,360});
+    // torch::Tensor output = net.forward({example}).toTensor();
+    // std::cout<<"ok"<<std::endl;
+}
\ No newline at end of file

cpp/utils/CMakeLists.txt

+# Set the project name
+project (utils)
+aux_source_directory(./src DIR_LIB_SRCS)
+
+add_library(${PROJECT_NAME} SHARED ${DIR_LIB_SRCS})
+
+# 添加.h文件查找路径
+target_include_directories( ${PROJECT_NAME}
+    PUBLIC ${PROJECT_SOURCE_DIR}/include
+)
+
+# 添加掉用的第三方库
+target_link_libraries(  "${PROJECT_NAME}" 
+                        "${OpenCV_LIBS}")
\ No newline at end of file

cpp/utils/include/data.hpp

+#ifndef DATA_H
+#define DATA_H
+#include <opencv2/opencv.hpp>
+
+namespace data {
+void normalize(cv::Mat& matrix, double mean = 0.5, double std = 0.5);
+
+}  // namespace data
+
+#endif
\ No newline at end of file

cpp/utils/include/util.hpp

+#ifndef UTIL_H
+#define UTIL_H
+#include <iostream>
+#include <list>
+namespace util {
+
+class Timer {
+   private:
+    clock_t tstart, tend;
+
+   public:
+    void start();
+    void end();
+};
+
+// std::string path = util::current_path();
+std::string current_path();
+
+// std::string out = util::pathjoin({path, "b", "c"});
+std::string pathjoin(const std::list<std::string>& strs);
+
+bool isfile(const std::string& name);
+
+}  // namespace util
+
+#endif
\ No newline at end of file

cpp/utils/src/data.cpp

+#include "data.hpp"
+#include <opencv2/opencv.hpp>
+
+namespace data {
+void normalize(cv::Mat& matrix, double mean, double std) {
+    // matrix = (matrix / 255.0 - mean) / std;
+    matrix = matrix / (255.0 * std) - mean / std;
+}
+
+}  // namespace data
\ No newline at end of file

cpp/utils/src/util.cpp

+#include "util.hpp"
+#include <stdio.h>
+#include <sys/stat.h>
+#include <unistd.h>
+#include <iostream>
+#include <list>
+#include <vector>
+
+namespace util {
+
+void Timer::start() {
+    tstart = clock();
+}
+void Timer::end() {
+    tend = clock();
+    double dur;
+    dur = (double)(tend - tstart);
+    std::cout << "Cost Time:" << (dur / CLOCKS_PER_SEC) << "\n";
+}
+
+std::string current_path() {
+    char* buffer;
+    buffer = getcwd(NULL, 0);
+    return buffer;
+}
+
+std::string pathjoin(const std::list<std::string>& strs) {
+    std::string res = "";
+    int cnt = 0;
+    for (std::string s : strs) {
+        if (cnt == 0) {
+            res += s;
+        } else {
+            if (s[0] != '/') {
+                res += ("/" + s);
+            } else {
+                res += s;
+            }
+        }
+        cnt++;
+    }
+    return res;
+}
+
+bool isfile(const std::string& name) {
+    struct stat buffer;
+    return (stat(name.c_str(), &buffer) == 0);
+}
+}  // namespace util
\ No newline at end of file

deepmosaic.py

@@ -2,7 +2,7 @@ import os
 import sys
 import traceback
 try:
-    from cores import Options,core
+    from cores import Options,add,clean,style
     from util import util
     from models import loadmodel
 except Exception as e:
@@ -25,9 +25,9 @@ def main():
         for file in files:
             opt.media_path = file
             if util.is_img(file):
-                core.addmosaic_img(opt,netS)
+                add.addmosaic_img(opt,netS)
             elif util.is_video(file):
-                core.addmosaic_video(opt,netS)
+                add.addmosaic_video(opt,netS)
                 util.clean_tempfiles(opt, tmp_init = False)
             else:
                 print('This type of file is not supported')
@@ -45,12 +45,12 @@ def main():
         for file in files:
             opt.media_path = file
             if util.is_img(file):
-                core.cleanmosaic_img(opt,netG,netM)
+                clean.cleanmosaic_img(opt,netG,netM)
             elif util.is_video(file):
                 if opt.netG == 'video' and not opt.traditional:            
-                    core.cleanmosaic_video_fusion(opt,netG,netM)
+                    clean.cleanmosaic_video_fusion(opt,netG,netM)
                 else:
-                    core.cleanmosaic_video_byframe(opt,netG,netM)
+                    clean.cleanmosaic_video_byframe(opt,netG,netM)
                 util.clean_tempfiles(opt, tmp_init = False)
             else:
                 print('This type of file is not supported')
@@ -60,9 +60,9 @@ def main():
         for file in files:
             opt.media_path = file
             if util.is_img(file):
-                core.styletransfer_img(opt,netG)
+                style.styletransfer_img(opt,netG)
             elif util.is_video(file):
-                core.styletransfer_video(opt,netG)
+                style.styletransfer_video(opt,netG)
                 util.clean_tempfiles(opt, tmp_init = False)
             else:
                 print('This type of file is not supported')
@@ -79,7 +79,7 @@ if __name__ == '__main__':
     except Exception as ex:
         print('--------------------ERROR--------------------')
         print('--------------Environment--------------')
-        print('DeepMosaics: 0.5.0')
+        print('DeepMosaics: 0.5.1')
         print('Python:',sys.version)
         import torch
         print('Pytorch:',torch.__version__)

docs/Release_notes.txt

-DeepMosaics V0.3.0
-Core program building with windows10_1703_x86_64 
- + python 3.68  
- + pyinstaller 3.5
+DeepMosaics: 0.5.1
+Core building with:
+    Python: 3.7.3 (default, Apr 24 2019, 15:29:51) [MSC v.1915 64 bit (AMD64)]
+    Pytorch: 1.7.1
+    OpenCV: 4.1.2
+    Platform: Windows-10-10.0.19041-SP0
+    Driver Version: 461.40
+    CUDA:11.0
 GUI building with C#
 For more detail, please view on github: https://github.com/HypoX64/DeepMosaics
 
 Releases History
- V0.3.0
-   1. Support BiSeNet(Better recognition of mosaics).
-   2. New videoHD model.
-   3. Better feathering method.
- V0.2.0
-   1. Add video model.
-   2. Now you can input chinese path
-   3. Support style transfer
-   4. Support fps limit
- V0.1.2
-   1. Support pix2pixHD model
- V0.1.1
-   1. Check path, can't input illegal path
- V0.1.0
-   1. Initial release.
\ No newline at end of file
+  V0.5.1
+    Fix:
+      1.Fix Some BUGs when restore unfinished tasks.
+      2.Fix that audio and video are not synchronized when the video is too long.
+    New:
+      1.Speed up video processing by Asynchronous.
+  V0.5.0
+    1.New video model (Perform better)
+  V0.4.1
+    1.Allow unfinished tasks to be restored.
+    2.Clean cache during processing.
+    3.Support CUDA 11.0.
+  V0.4.0
+    1.Support GPU.
+    2.Preview images when processing video.
+    3.Choose start position of video.
+  V0.3.0
+    1. Support BiSeNet(Better recognition of mosaics).
+    2. New videoHD model.
+    3. Better feathering method.
+  V0.2.0
+    1. Add video model.
+    2. Now you can input chinese path
+    3. Support style transfer
+    4. Support fps limit
+  V0.1.2
+    1. Support pix2pixHD model
+  V0.1.1
+    1. Check path, can't input illegal path
+  V0.1.0
+    1. Initial release.
\ No newline at end of file

server.py → tools/server.py

@@ -4,7 +4,7 @@ import traceback
 import cv2
 import numpy as np
 try:
-    from cores import Options,core
+    from cores import Options,clean
     from util import util
     from util import image_processing as impro
     from models import loadmodel
@@ -44,7 +44,7 @@ def handle():
     try:
         if max(img.shape)>1080:
             img = impro.resize(img,720,interpolation=cv2.INTER_CUBIC)
-        img = core.cleanmosaic_img_server(opt,img,netG,netM)
+        img = clean.cleanmosaic_img_server(opt,img,netG,netM)
     except Exception as e:
         result['img'] = imgRec
         result['info'] = 'procfailed'

tools/trace_model.py

+import os
+import sys
+import traceback
+sys.path.append("..")
+from util import mosaic
+import torch
+
+try:
+    from cores import Options,add,clean,style
+    from util import util
+    from models import loadmodel
+except Exception as e:
+    print(e)
+    input('Please press any key to exit.\n')
+    sys.exit(0)
+
+opt = Options().getparse(test_flag = False)
+if not os.path.isdir(opt.temp_dir):
+    util.file_init(opt)
+
+def saveScriptModel(model,example,savepath):
+    model.cpu()
+    traced_script_module = torch.jit.trace(model, example)
+    # try ScriptModel
+    output = traced_script_module(example)
+    print(output)
+    traced_script_module.save(savepath)
+
+savedir = '../cpp/res/models/'
+util.makedirs(savedir)
+
+opt.mosaic_position_model_path = '../pretrained_models/mosaic/mosaic_position.pth'
+model = loadmodel.bisenet(opt,'mosaic')
+example = torch.ones((1,3,360,360))
+saveScriptModel(model,example,os.path.join(savedir,'mosaic_position.pt'))
+
+
+
+# def main():
+    
+#     if os.path.isdir(opt.media_path):
+#         files = util.Traversal(opt.media_path)
+#     else:
+#         files = [opt.media_path]        
+#     if opt.mode == 'add':
+#         netS = loadmodel.bisenet(opt,'roi')
+#         for file in files:
+#             opt.media_path = file
+#             if util.is_img(file):
+#                 add.addmosaic_img(opt,netS)
+#             elif util.is_video(file):
+#                 add.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')
+#         if opt.traditional:
+#             netG = None
+#         elif opt.netG == 'video':
+#             netG = loadmodel.video(opt)
+#         else:
+#             netG = loadmodel.pix2pix(opt)
+        
+#         for file in files:
+#             opt.media_path = file
+#             if util.is_img(file):
+#                 clean.cleanmosaic_img(opt,netG,netM)
+#             elif util.is_video(file):
+#                 if opt.netG == 'video' and not opt.traditional:            
+#                     clean.cleanmosaic_video_fusion(opt,netG,netM)
+#                 else:
+#                     clean.cleanmosaic_video_byframe(opt,netG,netM)
+#                 util.clean_tempfiles(opt, tmp_init = False)
+#             else:
+#                 print('This type of file is not supported')
+
+#     elif opt.mode == 'style':
+#         netG = loadmodel.style(opt)
+#         for file in files:
+#             opt.media_path = file
+#             if util.is_img(file):
+#                 style.styletransfer_img(opt,netG)
+#             elif util.is_video(file):
+#                 style.styletransfer_video(opt,netG)
+#                 util.clean_tempfiles(opt, tmp_init = False)
+#             else:
+#                 print('This type of file is not supported')
+
+#     util.clean_tempfiles(opt, tmp_init = False)
+
+# if __name__ == '__main__':
+#     main()
\ No newline at end of file

util/ffmpeg.py

@@ -43,7 +43,7 @@ def video2image(videopath, imagepath, fps=0, start_time='00:00:00', last_time='0
     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']
+    args = ['ffmpeg', '-i', '"'+videopath+'"','-async 1 -f mp3','-b:a 320k']
     if last_time != '00:00:00':
         args += ['-ss', start_time]
         args += ['-t', last_time]

util/image_processing.py

 import cv2
 import numpy as np
 import random
+from threading import Thread
 
 import platform
 
@@ -39,15 +40,22 @@ def imread(file_path,mod = 'normal',loadsize = 0, rgb=False):
 
     return img
 
-def imwrite(file_path,img):
+def imwrite(file_path,img,use_thread=False):
     '''
     in other to save chinese path images in windows,
     this fun just for save final output images
     '''
-    if system_type == 'Linux':
-        cv2.imwrite(file_path, img)
+    def subfun(file_path,img):
+        if system_type == 'Linux':
+            cv2.imwrite(file_path, img)
+        else:
+            cv2.imencode('.jpg', img)[1].tofile(file_path)
+    if use_thread:
+        t = Thread(target=subfun,args=(file_path, img,))
+        t.daemon()
+        t.start
     else:
-        cv2.imencode('.jpg', img)[1].tofile(file_path)
+        subfun(file_path,img)
 
 def resize(img,size,interpolation=cv2.INTER_LINEAR):
     '''
@@ -183,7 +191,7 @@ def mask_area(mask):
     except:
         area = 0
     return area
-import time
+
 def replace_mosaic(img_origin,img_fake,mask,x,y,size,no_feather):
     img_fake = cv2.resize(img_fake,(size*2,size*2),interpolation=cv2.INTER_CUBIC)
     if no_feather:

util/util.py

+import json
 import os
 import random
 import string
@@ -52,13 +53,25 @@ def is_dirs(paths):
             tmp.append(path)
     return tmp  
 
-def  writelog(path,log,isprint=False):
+def writelog(path,log,isprint=False):
     f = open(path,'a+')
     f.write(log+'\n')
     f.close()
     if isprint:
         print(log)
 
+def savejson(path,data_dict):
+    json_str = json.dumps(data_dict)
+    f = open(path,'w+')
+    f.write(json_str)
+    f.close()
+
+def loadjson(path):
+    f = open(path, 'r')
+    txt_data = f.read()
+    f.close()
+    return json.loads(txt_data)
+
 def makedirs(path):
     if os.path.isdir(path):
         print(path,'existed')