diff --git a/.gitignore b/.gitignore index b39cd6f33913160cec40c4597edbe698b52071b0..03ab31855bc6f5fff72cb3a63b1e323593d9b219 100644 --- a/.gitignore +++ b/.gitignore @@ -145,6 +145,7 @@ result/ /pix2pix /pix2pixHD /tmp +/tmp_files /to_make_show /test_media /result diff --git a/cores/core.py b/cores/core.py index fe3f76f11dd7043f508686516c898791f5bc30f9..b320dbd1d0a3113a79f2f3fc788346c218a95470 100644 --- a/cores/core.py +++ b/cores/core.py @@ -21,7 +21,7 @@ def addmosaic_img(opt,netS): img = impro.imread(path) mask = runmodel.get_ROI_position(img,netS,opt)[0] img = mosaic.addmosaic(img,mask,opt) - cv2.imwrite(os.path.join(opt.result_dir,os.path.splitext(os.path.basename(path))[0]+'_add.jpg'),img) + 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 @@ -56,7 +56,7 @@ def styletransfer_img(opt,netG): img = impro.imread(opt.media_path) img = runmodel.run_styletransfer(opt, netG, img) suffix = os.path.basename(opt.model_path).replace('.pth','').replace('style_','') - cv2.imwrite(os.path.join(opt.result_dir,os.path.splitext(os.path.basename(opt.media_path))[0]+'_'+suffix+'.jpg'),img) + 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 @@ -89,7 +89,7 @@ def cleanmosaic_img(opt,netG,netM): img_result = impro.replace_mosaic(img_origin,img_fake,x,y,size,opt.no_feather) else: print('Do not find mosaic') - cv2.imwrite(os.path.join(opt.result_dir,os.path.splitext(os.path.basename(path))[0]+'_clean.jpg'),img_result) + impro.imwrite(os.path.join(opt.result_dir,os.path.splitext(os.path.basename(path))[0]+'_clean.jpg'),img_result) def cleanmosaic_video_byframe(opt,netG,netM): path = opt.media_path diff --git a/cores/options.py b/cores/options.py index 5f2264fd46a2d4c31e24f8d1600ff26c5c66c395..b621a026d2a0c813235ed35113f8700db3ed238d 100644 --- a/cores/options.py +++ b/cores/options.py @@ -17,9 +17,9 @@ class Options(): 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 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') + help='select model to use for netG(Clean mosaic and Transfer style) -> auto | unet_128 | unet_256 | resnet_9blocks | HD | video') + self.parser.add_argument('--output_size', type=int, default=0,help='size of output file,if 0 -> origin') #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') @@ -35,7 +35,8 @@ class Options(): 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('--preprocess', type=str, default='resize', help='resize and cropping of images at load time [ resize | resize_scale_width | edges | gray] or resize,edges(use comma to split)') + self.parser.add_argument('--edges', action='store_true', help='if true, use edges to generate pictures,(input_nc = 1)') 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') @@ -65,6 +66,12 @@ class Options(): else: print('Please input running mode!') + 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' @@ -77,12 +84,9 @@ 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 + self.opt.ex_mult = 1.1 else: self.opt.ex_mult = 1.5 else: diff --git a/deepmosaic.py b/deepmosaic.py index 3305689b94b22b42c144671ec8a5ffc0d6ffeee9..c75a06897dc9cd131fa31a02443682c898982c64 100644 --- a/deepmosaic.py +++ b/deepmosaic.py @@ -56,12 +56,12 @@ def main(): util.clean_tempfiles(tmp_init = False) -main() -# if __name__ == '__main__': -# try: -# main() -# except Exception as e: -# print('Error:',e) -# input('Please press any key to exit.\n') -# util.clean_tempfiles(tmp_init = False) -# exit(0) +# main() +if __name__ == '__main__': + try: + main() + except Exception as e: + print('Error:',e) + input('Please press any key to exit.\n') + util.clean_tempfiles(tmp_init = False) + exit(0) diff --git a/make_datasets/get_edges_pix2pix_dataset.py b/make_datasets/get_edges_pix2pix_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..5def0ebd91e28b1fdd74f9359acca1aa1ba1deb8 --- /dev/null +++ b/make_datasets/get_edges_pix2pix_dataset.py @@ -0,0 +1,26 @@ +import numpy as np +import cv2 +import os +import sys +sys.path.append("..") +from util import image_processing as impro +from util import util + +img_dir = './datasets_img/pix2pix/edges2cat/images' +output_dir = './datasets_img/pix2pix/edges2cat/train' +util.makedirs(output_dir) + +img_names = os.listdir(img_dir) +for i,img_name in enumerate(img_names,2000): + try: + img = impro.imread(os.path.join(img_dir,img_name)) + img = impro.resize(img, 286) + h,w = img.shape[:2] + edges = cv2.Canny(img,150,250) + edges = impro.ch_one2three(edges) + out_img = np.zeros((h,w*2,3), dtype=np.uint8) + out_img[:,0:w] = edges + out_img[:,w:2*w] = img + cv2.imwrite(os.path.join(output_dir,'%05d' % i+'.jpg'), out_img) + except Exception as e: + pass diff --git a/models/loadmodel.py b/models/loadmodel.py index 468ccd0ae00849abdfa07f9b852d9a34aed1620d..6257da075bd0919180395ce342609882bbb8bc36 100755 --- a/models/loadmodel.py +++ b/models/loadmodel.py @@ -39,7 +39,7 @@ def pix2pix(opt): def style(opt): if opt.edges: - netG = define_G(1, 3, 64, 'unet_256', norm='instance',use_dropout=True, init_type='normal', gpu_ids=[]) + netG = define_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=[]) diff --git a/models/runmodel.py b/models/runmodel.py index 2c7d4fcca07e77b5f4c3d129522556b144247c80..ff156f5a77ea83754a30dfd85e9674b9579e389e 100755 --- a/models/runmodel.py +++ b/models/runmodel.py @@ -37,11 +37,15 @@ def run_pix2pix(img,net,opt): return img_fake 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 'resize' in opt.preprocess and 'resize_scale_width' not in opt.preprocess: + img = impro.resize(img,opt.output_size) + elif 'resize_scale_width' in opt.preprocess: + img = cv2.resize(img, (opt.output_size,opt.output_size)) + img = img[0:4*int(img.shape[0]/4),0:4*int(img.shape[1]/4),:] + + if 'edges' in opt.preprocess: if opt.canny > 100: canny_low = opt.canny-50 canny_high = np.clip(opt.canny+50,0,255) @@ -56,7 +60,6 @@ def run_styletransfer(opt, net, img): 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) diff --git a/util/ffmpeg.py b/util/ffmpeg.py index 5ba4ab4ce08c53b80689e32fde4588c43b709432..9b02153e567779deb145910408628fb3feb2f217 100755 --- a/util/ffmpeg.py +++ b/util/ffmpeg.py @@ -45,4 +45,4 @@ def continuous_screenshot(videopath,savedir,fps): fps: save how many images per second ''' videoname = os.path.splitext(os.path.basename(videopath))[0] - os.system('ffmpeg -i '+videopath+' -vf fps='+str(fps)+' '+savedir+'/'+videoname+'_%05d.jpg') + os.system('ffmpeg -i "'+videopath+'" -vf fps='+str(fps)+' '+savedir+'/'+videoname+'_%05d.jpg') diff --git a/util/image_processing.py b/util/image_processing.py index 1a6f099c049ee3e2dfcc5899e68818f9283b65ee..56e7b7c580adebfc791f44c4955e441cc8e0722a 100755 --- a/util/image_processing.py +++ b/util/image_processing.py @@ -2,20 +2,42 @@ import cv2 import numpy as np import random +import platform +system_type = 'Linux' +if 'Windows' in platform.platform(): + system_type = 'Windows' + def imread(file_path,mod = 'normal'): ''' mod = 'normal' | 'gray' | 'all' ''' - if mod == 'normal': - cv_img = cv2.imread(file_path) - elif mod == 'gray': - cv_img = cv2.imread(file_path,0) - elif mod == 'all': - cv_img = cv2.imread(file_path,-1) - - # # imread for chinese path in windows but no EXIF - # cv_img = cv2.imdecode(np.fromfile(file_path,dtype=np.uint8),-1) - return cv_img + if system_type == 'Linux': + if mod == 'normal': + img = cv2.imread(file_path) + elif mod == 'gray': + img = cv2.imread(file_path,0) + elif mod == 'all': + img = cv2.imread(file_path,-1) + + #For chinese path, use cv2.imdecode in windows. + #It will loss EXIF, I can't fix it + else: + if mod == 'gray': + img = cv2.imdecode(np.fromfile(file_path,dtype=np.uint8),0) + else: + img = cv2.imdecode(np.fromfile(file_path,dtype=np.uint8),-1) + + return img + +def imwrite(file_path,img): + ''' + 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) + else: + cv2.imencode('.jpg', img)[1].tofile(file_path) def resize(img,size,interpolation=cv2.INTER_LINEAR): h, w = img.shape[:2]