allow edges2video

cores/options.py

@@ -13,23 +13,32 @@ def initialize(self):
         self.parser.add_argument('--use_gpu',type=int,default=1, help='if 0 or -1, do not use gpu')
         # self.parser.add_argument('--use_gpu', action='store_true', help='if input it, use gpu')
         self.parser.add_argument('--media_path', type=str, default='./hands_test.mp4',help='your videos or images path')
-        self.parser.add_argument('--mode', type=str, default='auto',help='add or clean mosaic into your media  auto | add | clean | style')
+        self.parser.add_argument('--mode', type=str, default='auto',help='auto | add | clean | style')
         self.parser.add_argument('--model_path', type=str, default='./pretrained_models/add_hands_128.pth',help='pretrained model path')
-        self.parser.add_argument('--result_dir', type=str, default='./result',help='output result will be saved here')
+        self.parser.add_argument('--result_dir', type=str, default='./result',help='output media will be saved here')
         self.parser.add_argument('--tempimage_type', type=str, default='png',help='type of temp image, png | jpg, png is better but occupy more storage space')
-
+        self.parser.add_argument('--output_size', type=int, default=0,help='size of output file,if 0 -> origin')
+        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')
+
         #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')
         self.parser.add_argument('--mosaic_size', type=int, default=0,help='mosaic size,if 0 auto size')
         self.parser.add_argument('--mask_extend', type=int, default=10,help='more mosaic area')
         self.parser.add_argument('--mask_threshold', type=int, default=64,help='threshold of recognize mosaic position 0~255')
-        self.parser.add_argument('--output_size', type=int, default=0,help='size of output file,if 0 -> origin')
 
-        #CleanMosaic
-        self.parser.add_argument('--netG', type=str, default='auto',help='select model to use for netG(clean mosaic) -> auto | unet_128 | resnet_9blocks | HD | video')
+        #CleanMosaic     
         self.parser.add_argument('--mosaic_position_model_path', type=str, default='auto',help='name of model use to find mosaic position')
         self.parser.add_argument('--no_feather', action='store_true', help='if true, no edge feather and color correction, but run faster')
+        self.parser.add_argument('--no_large_area', action='store_true', help='if true, do not find the largest mosaic area')
         self.parser.add_argument('--medfilt_num', type=int, default=11,help='medfilt window of mosaic movement in the video')
+        self.parser.add_argument('--ex_mult', type=str, default='auto',help='mosaic area expansion')
+
+        #StyleTransfer
+        self.parser.add_argument('--edges', action='store_true', help='if true, make edges first')
+        self.parser.add_argument('--canny', type=int, default=150,help='threshold of canny')
+        self.parser.add_argument('--only_edges', action='store_true', help='if true, output media will be edges')
+
         self.initialized = True
 
 
@@ -38,24 +47,25 @@ def getparse(self):
             self.initialize()
         self.opt = self.parser.parse_args()
 
+        model_name = os.path.basename(self.opt.model_path)
+
         if torch.cuda.is_available() and self.opt.use_gpu > 0:
             self.opt.use_gpu = True
         else:
             self.opt.use_gpu = False
 
 
         if self.opt.mode == 'auto':
-            if 'add' in self.opt.model_path:
+            if 'add' in model_name:
                 self.opt.mode = 'add'
-            elif 'clean' in self.opt.model_path:
+            elif 'clean' in model_name:
                 self.opt.mode = 'clean'
-            elif 'style' in self.opt.model_path:
+            elif 'style' in model_name or 'edges' in model_name:
                 self.opt.mode = 'style'
             else:
                 print('Please input running mode!')
 
         if self.opt.netG == 'auto' and self.opt.mode =='clean':
-            model_name = os.path.basename(self.opt.model_path)
             if 'unet_128' in model_name:
                 self.opt.netG = 'unet_128'
             elif 'resnet_9blocks' in model_name:
@@ -67,6 +77,16 @@ def getparse(self):
             else:
                 print('Type of Generator error!')
 
+        if 'edges' in model_name:
+            self.opt.edges = True
+
+        if self.opt.ex_mult == 'auto':
+            if 'face' in model_name:
+                self.opt.ex_mult = 1.2
+            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')

deepmosaic.py

@@ -43,7 +43,7 @@ def main():
                 print('This type of file is not supported')
 
     elif opt.mode == 'style':
-        netG = loadmodel.cyclegan(opt)
+        netG = loadmodel.style(opt)
         for file in files:
             opt.media_path = file
             if util.is_img(file):

models/loadmodel.py

@@ -37,9 +37,12 @@ def pix2pix(opt):
     return netG
 
 
-def cyclegan(opt):
-    netG = define_G(3, 3, 64, 'resnet_9blocks', norm='instance',use_dropout=False, init_type='normal', gpu_ids=[])
-
+def style(opt):
+    if opt.edges:
+        netG = define_G(1, 3, 64, 'unet_256', 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=[])
+
     #in other to load old pretrain model
     #https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/models/base_model.py
     if isinstance(netG, torch.nn.DataParallel):

models/runmodel.py

@@ -1,9 +1,11 @@
+import cv2
 import sys
 sys.path.append("..")
 import util.image_processing as impro
 from util import mosaic
 from util import data
 import torch
+import numpy as np
 
 def run_unet(img,net,size = 224,use_gpu = True):
     img=impro.image2folat(img,3)
@@ -34,11 +36,28 @@ def run_pix2pix(img,net,opt):
     img_fake = data.tensor2im(img_fake)
     return img_fake
 
-def run_styletransfer(opt, net, img, outsize = 720):
-    if min(img.shape[:2]) >= outsize:
-        img = impro.resize(img,outsize)
-    img = img[0:4*int(img.shape[0]/4),0:4*int(img.shape[1]/4),:]
-    img = data.im2tensor(img,use_gpu=opt.use_gpu)
+def run_styletransfer(opt, net, img):
+    if opt.output_size != 0:
+        img = impro.resize(img,opt.output_size)
+    if opt.edges:
+        if not opt.only_edges:
+            img = img[0:256*int(img.shape[0]/256),0:256*int(img.shape[1]/256),:]
+        if opt.canny > 100:
+            canny_low = opt.canny-50
+            canny_high = np.clip(opt.canny+50,0,255)
+        elif opt.canny < 50:
+            canny_low = np.clip(opt.canny-25,0,255)
+            canny_high = opt.canny+25
+        else:
+            canny_low = opt.canny-int(opt.canny/2)
+            canny_high = opt.canny+int(opt.canny/2)
+        img = cv2.Canny(img,opt.canny-50,opt.canny+50)
+        if opt.only_edges:
+            return img
+        img = data.im2tensor(img,use_gpu=opt.use_gpu,gray=True,use_transform = False,is0_1 = False)
+    else:    
+        img = img[0:4*int(img.shape[0]/4),0:4*int(img.shape[1]/4),:]
+        img = data.im2tensor(img,use_gpu=opt.use_gpu)
     img = net(img)
     img = data.tensor2im(img)
     return img
@@ -53,8 +72,9 @@ def get_mosaic_position(img_origin,net_mosaic_pos,opt,threshold = 128 ):
     mask = run_unet_rectim(img_origin,net_mosaic_pos,use_gpu = opt.use_gpu)
     #mask_1 = mask.copy()
     mask = impro.mask_threshold(mask,30,threshold)
-    mask = impro.find_best_ROI(mask)
-    x,y,size,area = impro.boundingSquare(mask,Ex_mul=1.5)
+    if not opt.no_large_area:
+        mask = impro.find_best_ROI(mask)
+    x,y,size,area = impro.boundingSquare(mask,Ex_mul=opt.ex_mult)
     rat = min(img_origin.shape[:2])/224.0
     x,y,size = int(rat*x),int(rat*y),int(rat*size)
     return x,y,size,mask

train/add/train.py

@@ -22,19 +22,19 @@
 
 LR = 0.0002
 EPOCHS = 100
-BATCHSIZE = 8
+BATCHSIZE = 16
 LOADSIZE = 256
 FINESIZE = 224
 CONTINUE = True
 use_gpu = True
 SAVE_FRE = 1
 MAX_LOAD = 30000
-#cudnn.benchmark = True
 
 
-dir_img = './datasets/mosaic/mosaic/'
-dir_mask = './datasets/mosaic/mask/'
-dir_checkpoint = 'checkpoints/mosaic/'
+
+dir_img = './datasets/face/origin_image/'
+dir_mask = './datasets/face/mask/'
+dir_checkpoint = 'checkpoints/face/'
 
 
 def Totensor(img,use_gpu=True):
@@ -115,6 +115,7 @@ def loadimage(dir_img,dir_mask,loadsize,eval_p):
     net.load_state_dict(torch.load(dir_checkpoint+'last.pth'))
 if use_gpu:
     net.cuda()
+    cudnn.benchmark = True
 
 
 optimizer = torch.optim.Adam(net.parameters(), lr=LR, betas=(0.9, 0.999))