From 27fffdb47ed6d4e0dd3cbb543f5ec77b7a5e932d Mon Sep 17 00:00:00 2001 From: hypox64 Date: Fri, 16 Aug 2019 16:48:44 +0800 Subject: [PATCH] correct color --- deepmosaic.py | 2 +- options.py | 12 ++++++++++-- util/image_processing.py | 6 +++++- 3 files changed, 16 insertions(+), 4 deletions(-) diff --git a/deepmosaic.py b/deepmosaic.py index 2787650..2c5bff7 100644 --- a/deepmosaic.py +++ b/deepmosaic.py @@ -68,7 +68,7 @@ def main(): img_result = img_origin.copy() if size != 0 : img_mosaic = img_origin[y-size:y+size,x-size:x+size] - img_fake=runmodel.run_pix2pix(img_mosaic,netG,opt) + img_fake = runmodel.run_pix2pix(img_mosaic,netG,opt) img_result = impro.replace_mosaic(img_origin,img_fake,x,y,size,opt.no_feather) cv2.imwrite(os.path.join(opt.result_dir,os.path.basename(path)),img_result) diff --git a/options.py b/options.py index d13928e..7860702 100644 --- a/options.py +++ b/options.py @@ -11,7 +11,7 @@ class Options(): #base 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='add',help='add or clean mosaic into your media add | clean') + self.parser.add_argument('--mode', type=str, default='auto',help='add or clean mosaic into your media auto | add | clean') 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('--tempimage_type', type=str, default='png',help='type of temp image, png | jpg, png is better but occupy more storage space') @@ -26,7 +26,7 @@ class Options(): #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') 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,but run faster') + 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('--medfilt_num', type=int, default=11,help='medfilt window of mosaic movement in the video') self.initialized = True @@ -36,6 +36,14 @@ class Options(): self.initialize() self.opt = self.parser.parse_args() + if self.opt.mode == 'auto': + if 'add' in self.opt.model_path: + self.opt.mode = 'add' + elif 'clean' in self.opt.model_path: + self.opt.mode = 'clean' + else: + print('Please input running mode!') + if self.opt.netG == 'auto' and self.opt.mode =='clean': if 'unet_128' in self.opt.model_path: self.opt.netG = 'unet_128' diff --git a/util/image_processing.py b/util/image_processing.py index 0bd1d38..c4ed2de 100755 --- a/util/image_processing.py +++ b/util/image_processing.py @@ -124,11 +124,15 @@ def mask_area(mask): def replace_mosaic(img_origin,img_fake,x,y,size,no_father): img_fake = resize(img_fake,size*2) - if no_father: img_origin[y-size:y+size,x-size:x+size]=img_fake img_result = img_origin else: + #color correction + RGB_origin = img_origin[y-size:y+size,x-size:x+size].mean(0).mean(0) + RGB_fake = img_fake.mean(0).mean(0) + for i in range(3):img_fake[:,:,i] = np.clip(img_fake[:,:,i]+RGB_origin[i]-RGB_fake[i],0,255) + #eclosion eclosion_num = int(size/5) entad = int(eclosion_num/2+2) mask = np.zeros(img_origin.shape, dtype='uint8') -- GitLab