diff --git a/cores/options.py b/cores/options.py index f4e54e28fb9648a249d066280178f24bd46fb640..5f2264fd46a2d4c31e24f8d1600ff26c5c66c395 100644 --- a/cores/options.py +++ b/cores/options.py @@ -13,23 +13,32 @@ class Options(): 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,6 +47,8 @@ class Options(): 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: @@ -45,17 +56,16 @@ class Options(): 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 @@ class Options(): 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') diff --git a/deepmosaic.py b/deepmosaic.py index ba2bda36a275edcb7a70ed673859185918706a15..3305689b94b22b42c144671ec8a5ffc0d6ffeee9 100644 --- a/deepmosaic.py +++ b/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): diff --git a/models/loadmodel.py b/models/loadmodel.py index b67d33e23c87afddf5370c04081363743d46f370..468ccd0ae00849abdfa07f9b852d9a34aed1620d 100755 --- a/models/loadmodel.py +++ b/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): diff --git a/models/runmodel.py b/models/runmodel.py index 605ce4324f5745e808b4f915e7ac02ae874d8a3b..2c7d4fcca07e77b5f4c3d129522556b144247c80 100755 --- a/models/runmodel.py +++ b/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 \ No newline at end of file diff --git a/train/add/train.py b/train/add/train.py index 5c0770255cb263b1393bfdfe05456b90990644d4..4d57814c33668e13d2977ce85b744fe3fc868d30 100644 --- a/train/add/train.py +++ b/train/add/train.py @@ -22,19 +22,19 @@ import torch.backends.cudnn as cudnn 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 @@ if CONTINUE: 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))