Completed the video core and faster feather

cores/core.py

@@ -244,11 +244,15 @@ def cleanmosaic_video_byframe(opt,netG,netM):
 
 def cleanmosaic_video_fusion(opt,netG,netM):
     path = opt.media_path
-    N = 25
-    if 'HD' in os.path.basename(opt.model_path):
-        INPUT_SIZE = 256
-    else:
-        INPUT_SIZE = 128
+    N,T,S = 2,5,3
+    LEFT_FRAME = (N*S)
+    POOL_NUM = LEFT_FRAME*2+1
+    INPUT_SIZE = 256
+    FRAME_POS = np.linspace(0, (T-1)*S,T,dtype=np.int64)
+    img_pool = []
+    previous_frame = None
+    init_flag = True
+
     fps,imagepaths,height,width = video_init(opt,path)
     positions = get_mosaic_positions(opt,netM,imagepaths,savemask=True)
     t1 = time.time()
@@ -258,37 +262,41 @@ def cleanmosaic_video_fusion(opt,netG,netM):
     # clean mosaic
     print('Clean Mosaic:')
     length = len(imagepaths)
-
-    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]
-
+        input_stream = []
         # 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.append(impro.imread(os.path.join(opt.temp_dir+'/video2image',imagepaths[np.clip(i+j-12,0,len(imagepaths)-1)])))
-        else:
+        if i==0 :# init
+            for j in range(POOL_NUM):
+                img_pool.append(impro.imread(os.path.join(opt.temp_dir+'/video2image',imagepaths[np.clip(i+j-LEFT_FRAME,0,len(imagepaths)-1)])))
+        else: # load next frame
             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_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()
 
         if size>100:
             try:#Avoid unknown errors
-                #reshape to network input shape
-                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,gpu_id=opt.gpu_id)
-                unmosaic_pred = netG(mosaic_input_tensor)              
-                img_fake = data.tensor2im(unmosaic_pred,rgb2bgr = False)
+                for pos in FRAME_POS:
+                    input_stream.append(impro.resize(img_pool[pos][y-size:y+size,x-size:x+size], INPUT_SIZE)[:,:,::-1])
+                if init_flag:
+                    init_flag = False
+                    previous_frame = input_stream[N]
+                    previous_frame = data.im2tensor(previous_frame,bgr2rgb=True,gpu_id=opt.gpu_id)
+
+                input_stream = np.array(input_stream).reshape(1,T,INPUT_SIZE,INPUT_SIZE,3).transpose((0,4,1,2,3))
+                input_stream = data.to_tensor(data.normalize(input_stream),gpu_id=opt.gpu_id)
+                unmosaic_pred = netG(input_stream,previous_frame)
+                img_fake = data.tensor2im(unmosaic_pred,rgb2bgr = True)
+                previous_frame = unmosaic_pred
+                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)
+                init_flag = True
+                print('Error:',e)
+        else:
+            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))      
 

cores/options.py

@@ -11,6 +11,7 @@ def __init__(self):
     def initialize(self):
 
         #base
+        self.parser.add_argument('--debug', action='store_true', help='if specified, start debug mode')
         self.parser.add_argument('--gpu_id', type=str,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')
@@ -58,8 +59,7 @@ def getparse(self, test_flag = False):
 
         model_name = os.path.basename(self.opt.model_path)
         self.opt.temp_dir = os.path.join(self.opt.temp_dir, 'DeepMosaics_temp')
-
-
+
         if self.opt.gpu_id != '-1':
             os.environ["CUDA_VISIBLE_DEVICES"] = str(self.opt.gpu_id)
             import torch
@@ -70,53 +70,59 @@ def getparse(self, test_flag = False):
 
         if test_flag:
             if not os.path.exists(self.opt.media_path):
-                print('Error: Bad media path!')
-                input('Please press any key to exit.\n')
-                sys.exit(0)
-
-        if self.opt.mode == 'auto':
-            if 'clean' in model_name or self.opt.traditional:
-                self.opt.mode = 'clean'
-            elif 'add' in model_name:
-                self.opt.mode = 'add'
-            elif 'style' in model_name or 'edges' in model_name:
-                self.opt.mode = 'style'
-            else:
-                print('Please input running model!')
+                print('Error: Media does not exist!')
                 input('Please press any key to exit.\n')
                 sys.exit(0)
-
-        if self.opt.output_size == 0 and self.opt.mode == 'style':
-            self.opt.output_size = 512
-
-        if 'edges' in model_name or 'edges' in self.opt.preprocess:
-            self.opt.edges = True
-
-        if self.opt.netG == 'auto' and self.opt.mode =='clean':
-            if 'unet_128' in model_name:
-                self.opt.netG = 'unet_128'
-            elif 'resnet_9blocks' in model_name:
-                self.opt.netG = 'resnet_9blocks'
-            elif 'HD' in model_name and 'video' not in model_name:
-                self.opt.netG = 'HD'
-            elif 'video' in model_name:
-                self.opt.netG = 'video'
-            else:
-                print('Type of Generator error!')
+            if not os.path.exists(self.opt.model_path):
+                print('Error: Model does not exist!')
                 input('Please press any key to exit.\n')
                 sys.exit(0)
 
-        if self.opt.ex_mult == 'auto':
-            if 'face' in model_name:
-                self.opt.ex_mult = 1.1
+            if self.opt.mode == 'auto':
+                if 'clean' in model_name or self.opt.traditional:
+                    self.opt.mode = 'clean'
+                elif 'add' in model_name:
+                    self.opt.mode = 'add'
+                elif 'style' in model_name or 'edges' in model_name:
+                    self.opt.mode = 'style'
+                else:
+                    print('Please check model_path!')
+                    input('Please press any key to exit.\n')
+                    sys.exit(0)
+
+            if self.opt.output_size == 0 and self.opt.mode == 'style':
+                self.opt.output_size = 512
+
+            if 'edges' in model_name or 'edges' in self.opt.preprocess:
+                self.opt.edges = True
+
+            if self.opt.netG == 'auto' and self.opt.mode =='clean':
+                if 'unet_128' in model_name:
+                    self.opt.netG = 'unet_128'
+                elif 'resnet_9blocks' in model_name:
+                    self.opt.netG = 'resnet_9blocks'
+                elif 'video' in model_name:
+                    self.opt.netG = 'video'
+                else:
+                    print('Type of Generator error!')
+                    input('Please press any key to exit.\n')
+                    sys.exit(0)
+
+            if self.opt.ex_mult == 'auto':
+                if 'face' in model_name:
+                    self.opt.ex_mult = 1.1
+                else:
+                    self.opt.ex_mult = 1.5
             else:
-                self.opt.ex_mult = 1.5
-        else:
-            self.opt.ex_mult = float(self.opt.ex_mult)
-
-        if self.opt.mosaic_position_model_path == 'auto':
-            _path = os.path.join(os.path.split(self.opt.model_path)[0],'mosaic_position.pth')
-            self.opt.mosaic_position_model_path = _path
-            # print(self.opt.mosaic_position_model_path)
+                self.opt.ex_mult = float(self.opt.ex_mult)
+
+            if self.opt.mosaic_position_model_path == 'auto':
+                _path = os.path.join(os.path.split(self.opt.model_path)[0],'mosaic_position.pth')
+                if os.path.isfile(_path):
+                    self.opt.mosaic_position_model_path = _path
+                else:
+                    input('Please check mosaic_position_model_path!')
+                    input('Please press any key to exit.\n')
+                    sys.exit(0)
 
         return self.opt

deepmosaic.py

@@ -68,8 +68,11 @@ def main():
                 print('This type of file is not supported')
 
     util.clean_tempfiles(opt, tmp_init = False)
-        
+
 if __name__ == '__main__':
+    if opt.debug:
+        main()
+        sys.exit(0)
     try:
         main()
         print('Finished!')

models/loadmodel.py

@@ -1,23 +1,23 @@
 import torch
 from . import model_util
-from .pix2pix_model import define_G
-from .pix2pixHD_model import define_G as define_G_HD
-from .video_model import MosaicNet
-from .videoHD_model import MosaicNet as MosaicNet_HD
+from .pix2pix_model import define_G as pix2pix_G
+from .pix2pixHD_model import define_G as pix2pixHD_G
+# from .video_model import MosaicNet
+# from .videoHD_model import MosaicNet as MosaicNet_HD
 from .BiSeNet_model import BiSeNet
+from .BVDNet import define_G as video_G
 
 def show_paramsnumber(net,netname='net'):
     parameters = sum(param.numel() for param in net.parameters())
     parameters = round(parameters/1e6,2)
     print(netname+' parameters: '+str(parameters)+'M')
 
-
 def pix2pix(opt):
     # print(opt.model_path,opt.netG)
     if opt.netG == 'HD':
-        netG = define_G_HD(3, 3, 64, 'global' ,4)
+        netG = pix2pixHD_G(3, 3, 64, 'global' ,4)
     else:
-        netG = define_G(3, 3, 64, opt.netG, norm='batch',use_dropout=True, init_type='normal', gpu_ids=[])
+        netG = pix2pix_G(3, 3, 64, opt.netG, norm='batch',use_dropout=True, init_type='normal', gpu_ids=[])
     show_paramsnumber(netG,'netG')
     netG.load_state_dict(torch.load(opt.model_path))
     netG = model_util.todevice(netG,opt.gpu_id)
@@ -27,9 +27,9 @@ def pix2pix(opt):
 
 def style(opt):
     if opt.edges:
-        netG = define_G(1, 3, 64, 'resnet_9blocks', norm='instance',use_dropout=True, init_type='normal', gpu_ids=[])
+        netG = pix2pix_G(1, 3, 64, 'resnet_9blocks', norm='instance',use_dropout=True, init_type='normal', gpu_ids=[])
     else:
-        netG = define_G(3, 3, 64, 'resnet_9blocks', norm='instance',use_dropout=False, init_type='normal', gpu_ids=[])
+        netG = pix2pix_G(3, 3, 64, 'resnet_9blocks', norm='instance',use_dropout=False, init_type='normal', gpu_ids=[])
 
     #in other to load old pretrain model
     #https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/models/base_model.py
@@ -51,10 +51,7 @@ def style(opt):
     return netG
 
 def video(opt):
-    if 'HD' in opt.model_path:
-        netG = MosaicNet_HD(3*25+1, 3, norm='instance')
-    else:
-        netG = MosaicNet(3*25+1, 3,norm = 'batch')
+    netG = video_G(N=2,n_blocks=1,gpu_id=opt.gpu_id)
     show_paramsnumber(netG,'netG')
     netG.load_state_dict(torch.load(opt.model_path))
     netG = model_util.todevice(netG,opt.gpu_id)