add out path control and save name prefix (#44)

README.md

@@ -97,7 +97,7 @@ You can put the testing images in the `inputs/TestWhole` folder. If you would li
 #### Testing on Face Restoration:
 [Note] If you want to compare CodeFormer in your paper, please run the following command indicating `--has_aligned` (for cropped and aligned face), as the command for the whole image will involve a process of face-background fusion that may damage hair texture on the boundary, which leads to unfair comparison.
 
-👨🏻 Face Restoration (cropped and aligned face)
+🧑🏻 Face Restoration (cropped and aligned face)
 ```
 # For cropped and aligned faces
 python inference_codeformer.py --w 0.5 --has_aligned --test_path [input folder]

app.py

@@ -103,98 +103,101 @@ os.makedirs('output', exist_ok=True)
 
 def inference(image, background_enhance, face_upsample, upscale, codeformer_fidelity):
     """Run a single prediction on the model"""
-    # take the default setting for the demo
-    has_aligned = False
-    only_center_face = False
-    draw_box = False
-    detection_model = "retinaface_resnet50"
-
-    upscale = int(upscale) # covert type to int
-    face_helper = FaceRestoreHelper(
-        upscale,
-        face_size=512,
-        crop_ratio=(1, 1),
-        det_model=detection_model,
-        save_ext="png",
-        use_parse=True,
-        device=device,
-    )
-    bg_upsampler = upsampler if background_enhance else None
-    face_upsampler = upsampler if face_upsample else None
-
-    img = cv2.imread(str(image), cv2.IMREAD_COLOR)
-
-    if has_aligned:
-        # the input faces are already cropped and aligned
-        img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR)
-        face_helper.is_gray = is_gray(img, threshold=5)
-        if face_helper.is_gray:
-            print('Grayscale input: True')
-        face_helper.cropped_faces = [img]
-    else:
-        face_helper.read_image(img)
-        # get face landmarks for each face
-        num_det_faces = face_helper.get_face_landmarks_5(
-          only_center_face=only_center_face, resize=640, eye_dist_threshold=5
-        )
-        print(f"\tdetect {num_det_faces} faces")
-        # align and warp each face
-        face_helper.align_warp_face()
-
-    # face restoration for each cropped face
-    for idx, cropped_face in enumerate(face_helper.cropped_faces):
-        # prepare data
-        cropped_face_t = img2tensor(
-            cropped_face / 255.0, bgr2rgb=True, float32=True
+    try: # global try
+        # take the default setting for the demo
+        has_aligned = False
+        only_center_face = False
+        draw_box = False
+        detection_model = "retinaface_resnet50"
+
+        upscale = int(upscale) # covert type to int
+        face_helper = FaceRestoreHelper(
+            upscale,
+            face_size=512,
+            crop_ratio=(1, 1),
+            det_model=detection_model,
+            save_ext="png",
+            use_parse=True,
+            device=device,
         )
-        normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
-        cropped_face_t = cropped_face_t.unsqueeze(0).to(device)
-
-        try:
-            with torch.no_grad():
-                output = codeformer_net(
-                    cropped_face_t, w=codeformer_fidelity, adain=True
-                )[0]
-                restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
-            del output
-            torch.cuda.empty_cache()
-        except Exception as error:
-            print(f"\tFailed inference for CodeFormer: {error}")
-            restored_face = tensor2img(
-                cropped_face_t, rgb2bgr=True, min_max=(-1, 1)
-            )
-
-        restored_face = restored_face.astype("uint8")
-        face_helper.add_restored_face(restored_face)
-
-    # paste_back
-    if not has_aligned:
-        # upsample the background
-        if bg_upsampler is not None:
-            # Now only support RealESRGAN for upsampling background
-            bg_img = bg_upsampler.enhance(img, outscale=upscale)[0]
+        bg_upsampler = upsampler if background_enhance else None
+        face_upsampler = upsampler if face_upsample else None
+
+        img = cv2.imread(str(image), cv2.IMREAD_COLOR)
+
+        if has_aligned:
+            # the input faces are already cropped and aligned
+            img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR)
+            face_helper.is_gray = is_gray(img, threshold=5)
+            if face_helper.is_gray:
+                print('Grayscale input: True')
+            face_helper.cropped_faces = [img]
         else:
-            bg_img = None
-        face_helper.get_inverse_affine(None)
-        # paste each restored face to the input image
-        if face_upsample and face_upsampler is not None:
-            restored_img = face_helper.paste_faces_to_input_image(
-                upsample_img=bg_img,
-                draw_box=draw_box,
-                face_upsampler=face_upsampler,
+            face_helper.read_image(img)
+            # get face landmarks for each face
+            num_det_faces = face_helper.get_face_landmarks_5(
+            only_center_face=only_center_face, resize=640, eye_dist_threshold=5
             )
-        else:
-            restored_img = face_helper.paste_faces_to_input_image(
-                upsample_img=bg_img, draw_box=draw_box
+            print(f"\tdetect {num_det_faces} faces")
+            # align and warp each face
+            face_helper.align_warp_face()
+
+        # face restoration for each cropped face
+        for idx, cropped_face in enumerate(face_helper.cropped_faces):
+            # prepare data
+            cropped_face_t = img2tensor(
+                cropped_face / 255.0, bgr2rgb=True, float32=True
             )
-
-    # save restored img
-    save_path = f'output/out.png'
-    imwrite(restored_img, str(save_path))
-
-    restored_img = cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB)
-    return restored_img, save_path
-
+            normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
+            cropped_face_t = cropped_face_t.unsqueeze(0).to(device)
+
+            try:
+                with torch.no_grad():
+                    output = codeformer_net(
+                        cropped_face_t, w=codeformer_fidelity, adain=True
+                    )[0]
+                    restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
+                del output
+                torch.cuda.empty_cache()
+            except Exception as error:
+                print(f"\tFailed inference for CodeFormer: {error}")
+                restored_face = tensor2img(
+                    cropped_face_t, rgb2bgr=True, min_max=(-1, 1)
+                )
+
+            restored_face = restored_face.astype("uint8")
+            face_helper.add_restored_face(restored_face)
+
+        # paste_back
+        if not has_aligned:
+            # upsample the background
+            if bg_upsampler is not None:
+                # Now only support RealESRGAN for upsampling background
+                bg_img = bg_upsampler.enhance(img, outscale=upscale)[0]
+            else:
+                bg_img = None
+            face_helper.get_inverse_affine(None)
+            # paste each restored face to the input image
+            if face_upsample and face_upsampler is not None:
+                restored_img = face_helper.paste_faces_to_input_image(
+                    upsample_img=bg_img,
+                    draw_box=draw_box,
+                    face_upsampler=face_upsampler,
+                )
+            else:
+                restored_img = face_helper.paste_faces_to_input_image(
+                    upsample_img=bg_img, draw_box=draw_box
+                )
+
+        # save restored img
+        save_path = f'output/out.png'
+        imwrite(restored_img, str(save_path))
+
+        restored_img = cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB)
+        return restored_img, save_path
+    except Exception as error:
+        print('global exception', error)
+        return None, None
 
 
 title = "CodeFormer: Robust Face Restoration and Enhancement Network"

inference_codeformer.py

@@ -52,9 +52,10 @@ if __name__ == '__main__':
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     parser = argparse.ArgumentParser()
 
-    parser.add_argument('--w', type=float, default=0.5, help='Balance the quality and fidelity')
-    parser.add_argument('--upscale', type=int, default=2, help='The final upsampling scale of the image. Default: 2')
-    parser.add_argument('--test_path', type=str, default='./inputs/cropped_faces')
+    parser.add_argument('-i', '--test_path', type=str, default='./inputs/cropped_faces')
+    parser.add_argument('-o', '--save_path', type=str, default=None)
+    parser.add_argument('-w', '--w', type=float, default=0.5, help='Balance the quality and fidelity')
+    parser.add_argument('-s', '--upscale', type=int, default=2, help='The final upsampling scale of the image. Default: 2')
     parser.add_argument('--has_aligned', action='store_true', help='Input are cropped and aligned faces')
     parser.add_argument('--only_center_face', action='store_true', help='Only restore the center face')
     # large det_model: 'YOLOv5l', 'retinaface_resnet50'
@@ -64,12 +65,14 @@ if __name__ == '__main__':
     parser.add_argument('--bg_upsampler', type=str, default='None', help='background upsampler. Optional: realesrgan')
     parser.add_argument('--face_upsample', action='store_true', help='face upsampler after enhancement.')
     parser.add_argument('--bg_tile', type=int, default=400, help='Tile size for background sampler. Default: 400')
+    parser.add_argument('--suffix', type=str, default=None, help='Suffix of the restored faces')
     parser.add_argument('--save_video_fps', type=int, default=24, help='frame rate for saving video. Default: 24')
 
     args = parser.parse_args()
 
     # ------------------------ input & output ------------------------
     w = args.w
+
     if args.test_path.endswith(('jpg', 'png')): # input single img path
         input_img_list = [args.test_path]
         result_root = f'results/test_img_{w}'
@@ -89,7 +92,10 @@ if __name__ == '__main__':
         # scan all the jpg and png images
         input_img_list = sorted(glob.glob(os.path.join(args.test_path, '*.[jp][pn]g')))
         result_root = f'results/{os.path.basename(args.test_path)}_{w}'
-    
+
+    if not args.save_path is None: # set output path
+        result_root = args.save_path
+
     test_img_num = len(input_img_list)
     # ------------------ set up background upsampler ------------------
     if args.bg_upsampler == 'realesrgan':
@@ -215,11 +221,15 @@ if __name__ == '__main__':
                 save_face_name = f'{basename}.png'
             else:
                 save_face_name = f'{basename}_{idx:02d}.png'
+            if args.suffix is not None:
+                save_face_name = f'{save_face_name[:-4]}_{args.suffix}.png'
             save_restore_path = os.path.join(result_root, 'restored_faces', save_face_name)
             imwrite(restored_face, save_restore_path)
 
         # save restored img
         if not args.has_aligned and restored_img is not None:
+            if args.suffix is not None:
+                basename = f'{basename}_{args.suffix}'
             save_restore_path = os.path.join(result_root, 'final_results', f'{basename}.png')
             imwrite(restored_img, save_restore_path)