add --alpha

yangxy · Jun 1, 2022 · c29e2c9 · c29e2c9
1 parent 242584d
commit c29e2c9
Show file tree

Hide file tree

Showing 19 changed files with 16 additions and 12 deletions.
diff --git a/demo.py b/demo.py
@@ -75,6 +75,7 @@ def generate_mask(H, W, img=None):
     parser.add_argument('--out_size', type=int, default=None, help='out resolution of GPEN')
     parser.add_argument('--channel_multiplier', type=int, default=2, help='channel multiplier of GPEN')
     parser.add_argument('--narrow', type=float, default=1, help='channel narrow scale')
+    parser.add_argument('--alpha', type=float, default=1, help='blending the results')
     parser.add_argument('--use_sr', action='store_true', help='use sr or not')
     parser.add_argument('--use_cuda', action='store_true', help='use cuda or not')
     parser.add_argument('--save_face', action='store_true', help='save face or not')
@@ -93,7 +94,7 @@ def generate_mask(H, W, img=None):
     os.makedirs(args.outdir, exist_ok=True)
 
     if args.task == 'FaceEnhancement': 
-        processer = FaceEnhancement(in_size=args.in_size, model=args.model, use_sr=args.use_sr, sr_model=args.sr_model, sr_scale=args.sr_scale, tile_size=args.tile_size, channel_multiplier=args.channel_multiplier, narrow=args.narrow, key=args.key, device='cuda' if args.use_cuda else 'cpu')
+        processer = FaceEnhancement(args, in_size=args.in_size, model=args.model, use_sr=args.use_sr, device='cuda' if args.use_cuda else 'cpu')
     elif args.task == 'FaceColorization':
         processer = FaceColorization(in_size=args.in_size, model=args.model, device='cuda' if args.use_cuda else 'cpu')
     elif args.task == 'FaceInpainting':

diff --git a/examples/outs-selfie/2020-12-10 175750_COMP.jpg b/examples/outs-selfie/2020-12-10 175750_COMP.jpg
diff --git a/examples/outs-selfie/2020-12-10 175750_GPEN.jpg b/examples/outs-selfie/2020-12-10 175750_GPEN.jpg
diff --git a/examples/outs-selfie/2020-12-10 175750_face00.jpg b/examples/outs-selfie/2020-12-10 175750_face00.jpg
diff --git a/examples/outs-selfie/2020-12-17 185036_COMP.jpg b/examples/outs-selfie/2020-12-17 185036_COMP.jpg
diff --git a/examples/outs-selfie/2020-12-17 185036_GPEN.jpg b/examples/outs-selfie/2020-12-17 185036_GPEN.jpg
diff --git a/examples/outs-selfie/2020-12-17 185036_face00.jpg b/examples/outs-selfie/2020-12-17 185036_face00.jpg
diff --git a/examples/outs-selfie/20200320.101220_COMP.jpg b/examples/outs-selfie/20200320.101220_COMP.jpg
diff --git a/examples/outs-selfie/20200320.101220_GPEN.jpg b/examples/outs-selfie/20200320.101220_GPEN.jpg
diff --git a/examples/outs-selfie/20200320.101220_face00.jpg b/examples/outs-selfie/20200320.101220_face00.jpg
diff --git a/examples/outs-selfie/2021-03-10 160843(29)_COMP.jpg b/examples/outs-selfie/2021-03-10 160843(29)_COMP.jpg
diff --git a/examples/outs-selfie/2021-03-10 160843(29)_GPEN.jpg b/examples/outs-selfie/2021-03-10 160843(29)_GPEN.jpg
diff --git a/examples/outs-selfie/2021-03-10 160843(29)_face00.jpg b/examples/outs-selfie/2021-03-10 160843(29)_face00.jpg
diff --git a/examples/outs-selfie/IMG20180716171517_COMP.jpg b/examples/outs-selfie/IMG20180716171517_COMP.jpg
diff --git a/examples/outs-selfie/IMG20180716171517_GPEN.jpg b/examples/outs-selfie/IMG20180716171517_GPEN.jpg
diff --git a/examples/outs-selfie/IMG20180716171517_face00.jpg b/examples/outs-selfie/IMG20180716171517_face00.jpg
diff --git a/face_enhancement.py b/face_enhancement.py
@@ -13,21 +13,22 @@
 from align_faces import warp_and_crop_face, get_reference_facial_points
 
 class FaceEnhancement(object):
-    def __init__(self, base_dir='./', in_size=512, out_size=None, model=None, use_sr=True, sr_model=None, sr_scale=2, tile_size=0, channel_multiplier=2, narrow=1, key=None, device='cuda'):
+    def __init__(self, args, base_dir='./', in_size=512, out_size=None, model=None, use_sr=True, device='cuda'):
         self.facedetector = RetinaFaceDetection(base_dir, device)
-        self.facegan = FaceGAN(base_dir, in_size, out_size, model, channel_multiplier, narrow, key, device=device)
-        self.srmodel =  RealESRNet(base_dir, sr_model, sr_scale, tile_size, device=device)
+        self.facegan = FaceGAN(base_dir, in_size, out_size, model, args.channel_multiplier, args.narrow, args.key, device=device)
+        self.srmodel =  RealESRNet(base_dir, args.sr_model, args.sr_scale, args.tile_size, device=device)
         self.faceparser = FaceParse(base_dir, device=device)
         self.use_sr = use_sr
         self.in_size = in_size
         self.out_size = in_size if out_size is None else out_size
         self.threshold = 0.9
+        self.alpha = args.alpha
 
         # the mask for pasting restored faces back
         self.mask = np.zeros((512, 512), np.float32)
         cv2.rectangle(self.mask, (26, 26), (486, 486), (1, 1, 1), -1, cv2.LINE_AA)
-        self.mask = cv2.GaussianBlur(self.mask, (101, 101), 11)
-        self.mask = cv2.GaussianBlur(self.mask, (101, 101), 11)
+        self.mask = cv2.GaussianBlur(self.mask, (101, 101), 4)
+        self.mask = cv2.GaussianBlur(self.mask, (101, 101), 4)
 
         self.kernel = np.array((
                 [0.0625, 0.125, 0.0625],
@@ -41,11 +42,11 @@ def __init__(self, base_dir='./', in_size=512, out_size=None, model=None, use_sr
         self.reference_5pts = get_reference_facial_points(
                 (self.in_size, self.in_size), inner_padding_factor, outer_padding, default_square)
 
-    def mask_postprocess(self, mask, thres=20):
+    def mask_postprocess(self, mask, thres=26):
         mask[:thres, :] = 0; mask[-thres:, :] = 0
         mask[:, :thres] = 0; mask[:, -thres:] = 0
-        mask = cv2.GaussianBlur(mask, (101, 101), 11)
-        mask = cv2.GaussianBlur(mask, (101, 101), 11)
+        mask = cv2.GaussianBlur(mask, (101, 101), 4)
+        mask = cv2.GaussianBlur(mask, (101, 101), 4)
         return mask.astype(np.float32)
 
     def process(self, img, aligned=False):
@@ -92,6 +93,8 @@ def process(self, img, aligned=False):
 
             if min(fh, fw)<100: # gaussian filter for small faces
                 ef = cv2.filter2D(ef, -1, self.kernel)
+
+            ef = cv2.addWeighted(ef, self.alpha, of, 1.-self.alpha, 0.0)
 
             if self.in_size!=self.out_size:
                 ef = cv2.resize(ef, (self.in_size, self.in_size))

diff --git a/run.sh b/run.sh
@@ -1,9 +1,9 @@
 ############## TEST ################
 #python demo.py --task FaceEnhancement --model GPEN-BFR-512 --in_size 512 --channel_multiplier 2 --narrow 1 --use_sr --sr_scale 2 --use_cuda --save_face --indir examples/imgs --outdir examples/outs-bfr
 #python demo.py --task FaceEnhancement --model GPEN-BFR-256 --in_size 256 --channel_multiplier 1 --narrow 0.5 --use_sr --sr_scale 4 --use_cuda --save_face --indir examples/imgs --outdir examples/outs-bfr
-python demo.py --task FaceEnhancement --model GPEN-BFR-2048 --in_size 2048 --channel_multiplier 2 --narrow 1 --use_sr --sr_scale 4 --use_cuda --tile_size 400 --save_face --indir examples/selfie --outdir examples/outs-selfie
+python demo.py --task FaceEnhancement --model GPEN-BFR-2048 --in_size 2048 --channel_multiplier 2 --narrow 1 --alpha 0.8 --use_sr --sr_scale 2 --use_cuda --tile_size 400 --save_face --indir examples/selfie --outdir examples/outs-selfie
 
-#python demo.py --task FaceColorization --model GPEN-Colorization-1024 --in_size 1024 --use_cuda --indir examples/grays --outdir examples/outs-colorization
+#python demo.py --task FaceColorization --model GPEN-Colorization-1024 --in_size 1024 --use_cuda --indir examples/tmp --outdir examples/outs-colorization
 
 #python demo.py --task FaceInpainting --model GPEN-Inpainting-1024 --in_size 1024 --use_cuda --indir examples/ffhq-10 --outdir examples/outs-inpainting
 

diff --git a/sr_model/real_esrnet.py b/sr_model/real_esrnet.py
@@ -17,7 +17,7 @@ def __init__(self, base_dir='./', model=None, scale=2, tile_size=0, tile_pad=10,
     def load_srmodel(self, base_dir, model):
         self.srmodel = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=32, num_block=23, num_grow_ch=32, scale=self.scale)
         if model is None:
-            loadnet = torch.load(os.path.join(self.base_dir, 'weights', 'realesrnet_x2.pth'))
+            loadnet = torch.load(os.path.join(self.base_dir, 'weights', 'realesrnet_x%d.pth'%self.scale))
         else:
             loadnet = torch.load(os.path.join(self.base_dir, 'weights', model+'_x%d.pth'%self.scale))
         #print(loadnet['params_ema'].keys)