1024

inference.py

@@ -33,7 +33,7 @@ def main():
     print('[Info] Using {} for inference.'.format(device))
     os.makedirs(os.path.join('temp', args.tmp_dir), exist_ok=True)
 
-    enhancer = FaceEnhancement(base_dir='checkpoints', size=512, model='GPEN-BFR-512', use_sr=False, \
+    enhancer = FaceEnhancement(base_dir='checkpoints', size=1024, model='GPEN-BFR-1024', use_sr=False, \
                                sr_model='rrdb_realesrnet_psnr', channel_multiplier=2, narrow=1, device=device)
     restorer = GFPGANer(model_path='checkpoints/GFPGANv1.3.pth', upscale=1, arch='clean', \
                         channel_multiplier=2, bg_upsampler=None)
@@ -149,7 +149,7 @@ def main():
 
     if not os.path.isfile('temp/'+base_name+'_stablized.npy') or args.re_preprocess:
         imgs = []
-        for idx in tqdm(range(len(frames_pil)), desc="[Step 3] Stablize the expression In Video:"):
+        for idx in tqdm(range(len(frames_pil)), desc="[Step 3] Stabilize the expression In Video:"):
             if args.one_shot:
                 source_img = trans_image(frames_pil[0]).unsqueeze(0).to(device)
                 semantic_source_numpy = semantic_npy[0:1]
@@ -168,7 +168,7 @@ def main():
         np.save('temp/'+base_name+'_stablized.npy',imgs)
         del D_Net
     else:
-        print('[Step 3] Using saved stablized video.')
+        print('[Step 3] Using saved stabilized video.')
         imgs = np.load('temp/'+base_name+'_stablized.npy')
     torch.cuda.empty_cache()
 
@@ -249,22 +249,24 @@ def main():
             
             ff = xf.copy() 
             ff[y1:y2, x1:x2] = p
-            
+            height, width = ff.shape[:2]
+            pp = np.uint8(cv2.resize(np.clip(ff, 0 ,512), (width, height)))
+
+            pp, orig_faces, enhanced_faces = enhancer.process(pp, xf, bbox=c, face_enhance=True, possion_blending=False)
             # month region enhancement by GFPGAN
             cropped_faces, restored_faces, restored_img = restorer.enhance(
-                ff, has_aligned=False, only_center_face=True, paste_back=True)
+                pp, has_aligned=False, only_center_face=True, paste_back=True)
                 # 0,   1,   2,   3,   4,   5,   6,   7,   8,  9, 10,  11,  12,
             mm = [0,   0,   0,   0,   0,   0,   0,   0,   0,  0, 255, 255, 255, 0, 0, 0, 0, 0, 0]
+            #mm = [0, 255, 255, 255, 255, 255, 255, 255, 0, 0, 255, 255, 255, 0, 0, 0, 0, 0, 0]
             mouse_mask = np.zeros_like(restored_img)
             tmp_mask = enhancer.faceparser.process(restored_img[y1:y2, x1:x2], mm)[0]
             mouse_mask[y1:y2, x1:x2]= cv2.resize(tmp_mask, (x2 - x1, y2 - y1))[:, :, np.newaxis] / 255.
 
-            height, width = ff.shape[:2]
-            restored_img, ff, full_mask = [cv2.resize(x, (512, 512)) for x in (restored_img, ff, np.float32(mouse_mask))]
+           
+            restored_img, ff, full_mask = [cv2.resize(x, (1024, 1024)) for x in (restored_img, ff, np.float32(mouse_mask))]
             img = Laplacian_Pyramid_Blending_with_mask(restored_img, ff, full_mask[:, :, 0], 10)
-            pp = np.uint8(cv2.resize(np.clip(img, 0 ,255), (width, height)))
-
-            pp, orig_faces, enhanced_faces = enhancer.process(pp, xf, bbox=c, face_enhance=False, possion_blending=True)
+            pp = np.uint8(cv2.resize(np.clip(img, 0 ,1024), (width, height)))
             out.write(pp)
     out.release()
     
@@ -342,4 +344,4 @@ def datagen(frames, mels, full_frames, frames_pil, cox):
 
 
 if __name__ == '__main__':
-    main()
+    main()