add: app.py

app.py

@@ -0,0 +1,470 @@
+import os
+import uuid
+import glob
+import shutil
+from pathlib import Path
+from multiprocessing.pool import Pool
+
+import gradio as gr
+import torch
+from torchvision import transforms
+
+import cv2
+import numpy as np
+from PIL import Image
+import tqdm
+
+# from modules.networks.faceshifter import FSGenerator
+# from inference.alignment import norm_crop, norm_crop_with_M, paste_back
+# from inference.utils import save, get_5_from_98, get_detector, get_lmk
+# from inference.PIPNet.lib.tools import get_lmk_model, demo_image
+# from inference.landmark_smooth import kalman_filter_landmark, savgol_filter_landmark
+# from tricks import Trick
+
+# make_abs_path = lambda fn: os.path.abspath(os.path.join(os.path.dirname(os.path.realpath(__file__)), fn))
+#
+#
+# fs_model_name = 'faceshifter'
+# in_size = 512
+#
+# mouth_net_param = {
+#     "use": True,
+#     "feature_dim": 128,
+#     "crop_param": (28, 56, 84, 112),
+#     "weight_path": "../../modules/third_party/arcface/weights/mouth_net_28_56_84_112.pth",
+# }
+# trick = Trick()
+#
+# T = transforms.Compose(
+#         [
+#             transforms.ToTensor(),
+#             transforms.Normalize(0.5, 0.5),
+#         ]
+#     )
+# tensor2pil_transform = transforms.ToPILImage()
+#
+#
+# def extract_generator(ckpt: str, pt: str):
+#     print(f'[extract_generator] loading ckpt...')
+#     from trainer.faceshifter.faceshifter_pl import FaceshifterPL512, FaceshifterPL
+#     import yaml
+#     with open(make_abs_path('../../trainer/faceshifter/config.yaml'), 'r') as f:
+#         config = yaml.load(f, Loader=yaml.FullLoader)
+#     config['mouth_net'] = mouth_net_param
+#
+#     if in_size == 256:
+#         net = FaceshifterPL(n_layers=3, num_D=3, config=config)
+#     elif in_size == 512:
+#         net = FaceshifterPL512(n_layers=3, num_D=3, config=config, verbose=False)
+#     else:
+#         raise ValueError('Not supported in_size.')
+#     checkpoint = torch.load(ckpt, map_location="cpu", )
+#     net.load_state_dict(checkpoint["state_dict"], strict=False)
+#     net.eval()
+#
+#     G = net.generator
+#     torch.save(G.state_dict(), pt)
+#     print(f'[extract_generator] extracted from {ckpt}, pth saved to {pt}')
+#
+#
+# ''' load model '''
+# if fs_model_name == 'faceshifter':
+#     # pt_path = make_abs_path("../ffplus/extracted_ckpt/G_mouth1_t38.pth")
+#     # pt_path = make_abs_path("../ffplus/extracted_ckpt/G_mouth1_t512_6.pth")
+#     # ckpt_path = "/apdcephfs/share_1290939/gavinyuan/out/triplet512_6/epoch=3-step=128999.ckpt"
+#     pt_path = make_abs_path("../ffplus/extracted_ckpt/G_mouth1_t512_4.pth")
+#     ckpt_path = "/apdcephfs/share_1290939/gavinyuan/out/triplet512_4/epoch=2-step=185999.ckpt"
+#     if not os.path.exists(pt_path) or 't512' in pt_path:
+#         extract_generator(ckpt_path, pt_path)
+#     fs_model = FSGenerator(
+#         make_abs_path("../../modules/third_party/arcface/weights/ms1mv3_arcface_r100_fp16/backbone.pth"),
+#         mouth_net_param=mouth_net_param,
+#         in_size=in_size,
+#         downup=in_size == 512,
+#     )
+#     fs_model.load_state_dict(torch.load(pt_path, "cpu"), strict=True)
+#     fs_model.eval()
+#
+#     @torch.no_grad()
+#     def infer_batch_to_img(i_s, i_t, post: bool = False):
+#         i_r = fs_model(i_s, i_t)[0]  # x, id_vector, att
+#
+#         if post:
+#             target_hair_mask = trick.get_any_mask(i_t, par=[0, 17])
+#             target_hair_mask = trick.smooth_mask(target_hair_mask)
+#             i_r = target_hair_mask * i_t + (target_hair_mask * (-1) + 1) * i_r
+#             i_r = trick.finetune_mouth(i_s, i_t, i_r) if in_size == 256 else i_r
+#
+#         img_r = trick.tensor_to_arr(i_r)[0]
+#         return img_r
+#
+# elif fs_model_name == 'simswap_triplet' or fs_model_name == 'simswap_vanilla':
+#     from modules.networks.simswap import Generator_Adain_Upsample
+#     sw_model = Generator_Adain_Upsample(
+#         input_nc=3, output_nc=3, latent_size=512, n_blocks=9, deep=False,
+#         mouth_net_param=mouth_net_param
+#     )
+#     if fs_model_name == 'simswap_triplet':
+#         pt_path = make_abs_path("../ffplus/extracted_ckpt/G_mouth1_st5.pth")
+#         ckpt_path = make_abs_path("/apdcephfs/share_1290939/gavinyuan/out/"
+#                                   "simswap_triplet_5/epoch=12-step=782999.ckpt")
+#     elif fs_model_name == 'simswap_vanilla':
+#         pt_path = make_abs_path("../ffplus/extracted_ckpt/G_tmp_sv4_off.pth")
+#         ckpt_path = make_abs_path("/apdcephfs/share_1290939/gavinyuan/out/"
+#                                   "simswap_vanilla_4/epoch=694-step=1487999.ckpt")
+#     else:
+#         pt_path = None
+#         ckpt_path = None
+#     sw_model.load_state_dict(torch.load(pt_path, "cpu"), strict=False)
+#     sw_model.eval()
+#     fs_model = sw_model
+#
+#     from trainer.simswap.simswap_pl import SimSwapPL
+#     import yaml
+#     with open(make_abs_path('../../trainer/simswap/config.yaml'), 'r') as f:
+#         config = yaml.load(f, Loader=yaml.FullLoader)
+#     config['mouth_net'] = mouth_net_param
+#     net = SimSwapPL(config=config, use_official_arc='off' in pt_path)
+#
+#     checkpoint = torch.load(ckpt_path, map_location="cpu")
+#     net.load_state_dict(checkpoint["state_dict"], strict=False)
+#     net.eval()
+#     sw_mouth_net = net.mouth_net  # maybe None
+#     sw_netArc = net.netArc
+#     fs_model = fs_model.cuda()
+#     sw_mouth_net = sw_mouth_net.cuda() if sw_mouth_net is not None else sw_mouth_net
+#     sw_netArc = sw_netArc.cuda()
+#
+#     @torch.no_grad()
+#     def infer_batch_to_img(i_s, i_t, post: bool = False):
+#         i_r = fs_model(source=i_s, target=i_t, net_arc=sw_netArc, mouth_net=sw_mouth_net,)
+#         if post:
+#             target_hair_mask = trick.get_any_mask(i_t, par=[0, 17])
+#             target_hair_mask = trick.smooth_mask(target_hair_mask)
+#             i_r = target_hair_mask * i_t + (target_hair_mask * (-1) + 1) * i_r
+#         i_r = i_r.clamp(-1, 1)
+#         i_r = trick.tensor_to_arr(i_r)[0]
+#         return i_r
+#
+# elif fs_model_name == 'simswap_official':
+#     from simswap.image_infer import SimSwapOfficialImageInfer
+#     fs_model = SimSwapOfficialImageInfer()
+#     pt_path = 'Simswap Official'
+#     mouth_net_param = {
+#         "use": False
+#     }
+#
+#     @torch.no_grad()
+#     def infer_batch_to_img(i_s, i_t):
+#         i_r = fs_model.image_infer(source_tensor=i_s, target_tensor=i_t)
+#         i_r = i_r.clamp(-1, 1)
+#         return i_r
+#
+# else:
+#     raise ValueError('Not supported fs_model_name.')
+#
+#
+# print(f'[demo] model loaded from {pt_path}')
+
+
+def swap_image(
+    source_image,
+    target_path,
+    out_path,
+    transform,
+    G,
+    align_source="arcface",
+    align_target="set1",
+    gpu_mode=True,
+    paste_back=True,
+    use_post=False,
+    use_gpen=False,
+    in_size=256,
+):
+    name = target_path.split("/")[-1]
+    name = "out_" + name
+    if isinstance(G, torch.nn.Module):
+        G.eval()
+        if gpu_mode:
+            G = G.cuda()
+    source_img = np.array(Image.open(source_image).convert("RGB"))
+    net, detector = get_lmk_model()
+    lmk = get_5_from_98(demo_image(source_img, net, detector)[0])
+    source_img = norm_crop(source_img, lmk, in_size, mode=align_source, borderValue=0.0)
+    source_img = transform(source_img).unsqueeze(0)
+
+    target = np.array(Image.open(target_path).convert("RGB"))
+    original_target = target.copy()
+    lmk = get_5_from_98(demo_image(target, net, detector)[0])
+    target, M = norm_crop_with_M(target, lmk, in_size, mode=align_target, borderValue=0.0)
+    target = transform(target).unsqueeze(0)
+    if gpu_mode:
+        target = target.cuda()
+        source_img = source_img.cuda()
+
+    cv2.imwrite('cropped_source.png', trick.tensor_to_arr(source_img)[0, :, :, ::-1])
+    cv2.imwrite('cropped_target.png', trick.tensor_to_arr(target)[0, :, :, ::-1])
+
+    # both inputs should be 512
+    result = infer_batch_to_img(source_img, target, post=use_post)
+
+    cv2.imwrite('result.png', result[:, :, ::-1])
+
+    os.makedirs(out_path, exist_ok=True)
+    Image.fromarray(result.astype(np.uint8)).save(os.path.join(out_path, name))
+    save((result, M, original_target, os.path.join(out_path, "paste_back_" + name), None),
+         trick=trick, use_post=use_gpen)
+
+
+def process_video(
+    source_image,
+    target_path,
+    out_path,
+    transform,
+    G,
+    align_source="arcface",
+    align_target="set1",
+    gpu_mode=True,
+    frames=9999999,
+    use_tddfav2=False,
+    landmark_smooth="kalman",
+):
+    if isinstance(G, torch.nn.Module):
+        G.eval()
+        if gpu_mode:
+            G = G.cuda()
+    ''' Target video to frames (.png) '''
+    fps = 25.0
+    if not os.path.isdir(target_path):
+        vidcap = cv2.VideoCapture(target_path)
+        fps = vidcap.get(cv2.CAP_PROP_FPS)
+        try:
+            for match in glob.glob(os.path.join("./tmp/", "*.png")):
+                os.remove(match)
+            for match in glob.glob(os.path.join(out_path, "*.png")):
+                os.remove(match)
+        except Exception as e:
+            print(e)
+        os.makedirs("./tmp/", exist_ok=True)
+        os.system(
+            f"ffmpeg -i {target_path} -qscale:v 1 -qmin 1 -qmax 1 -vsync 0  ./tmp/frame_%05d.png"
+        )
+        target_path = "./tmp/"
+    globbed_images = sorted(glob.glob(os.path.join(target_path, "*.png")))
+    ''' Get target landmarks '''
+    print('[Extracting target landmarks...]')
+    if not use_tddfav2:
+        align_net, align_detector = get_lmk_model()
+    else:
+        align_net, align_detector = get_detector(gpu_mode=gpu_mode)
+    target_lmks = []
+    for frame_path in tqdm.tqdm(globbed_images):
+        target = np.array(Image.open(frame_path).convert("RGB"))
+        lmk = demo_image(target, align_net, align_detector)
+        lmk = lmk[0]
+        target_lmks.append(lmk)
+    ''' Landmark smoothing '''
+    target_lmks = np.array(target_lmks, np.float32)  # (#frames, 98, 2)
+    if landmark_smooth == 'kalman':
+        target_lmks = kalman_filter_landmark(target_lmks,
+                                             process_noise=0.01,
+                                             measure_noise=0.01).astype(np.int)
+    elif landmark_smooth == 'savgol':
+        target_lmks = savgol_filter_landmark(target_lmks).astype(np.int)
+    elif landmark_smooth == 'cancel':
+        target_lmks = target_lmks.astype(np.int)
+    else:
+        raise KeyError('Not supported landmark_smooth choice')
+    ''' Crop source image '''
+    source_img = np.array(Image.open(source_image).convert("RGB"))
+    if not use_tddfav2:
+        lmk = get_5_from_98(demo_image(source_img, align_net, align_detector)[0])
+    else:
+        lmk = get_lmk(source_img, align_net, align_detector)
+    source_img = norm_crop(source_img, lmk, in_size, mode=align_source, borderValue=0.0)
+    source_img = transform(source_img).unsqueeze(0)
+    if gpu_mode:
+        source_img = source_img.cuda()
+    ''' Process by frames '''
+    targets = []
+    t_facial_masks = []
+    Ms = []
+    original_frames = []
+    names = []
+    count = 0
+    for image in tqdm.tqdm(globbed_images):
+        names.append(os.path.join(out_path, Path(image).name))
+        target = np.array(Image.open(image).convert("RGB"))
+        original_frames.append(target)
+        ''' Crop target frames '''
+        lmk = get_5_from_98(target_lmks[count])
+        target, M = norm_crop_with_M(target, lmk, in_size, mode=align_target, borderValue=0.0)
+        target = transform(target).unsqueeze(0)  # in [-1,1]
+        if gpu_mode:
+            target = target.cuda()
+        ''' Finetune paste masks '''
+        target_facial_mask = trick.get_any_mask(target,
+                                                par=[1, 2, 3, 4, 5, 6, 10, 11, 12, 13]).squeeze()  # in [0,1]
+        target_facial_mask = target_facial_mask.cpu().numpy().astype(np.float)
+        target_facial_mask = trick.finetune_mask(target_facial_mask, target_lmks)  # in [0,1]
+        t_facial_masks.append(target_facial_mask)
+        ''' Face swapping '''
+        with torch.no_grad():
+            if 'faceshifter' in fs_model_name:
+                output = G(source_img, target)
+                target_hair_mask = trick.get_any_mask(target, par=[0, 17])
+                target_hair_mask = trick.smooth_mask(target_hair_mask)
+                output = target_hair_mask * target + (target_hair_mask * (-1) + 1) * output
+                output = trick.finetune_mouth(source_img, target, output)
+            elif 'simswap' in fs_model_name and 'official' not in fs_model_name:
+                output = fs_model(source=source_img, target=target,
+                                  net_arc=sw_netArc, mouth_net=sw_mouth_net,)
+                if 'vanilla' not in fs_model_name:
+                    target_hair_mask = trick.get_any_mask(target, par=[0, 17])
+                    target_hair_mask = trick.smooth_mask(target_hair_mask)
+                    output = target_hair_mask * target + (target_hair_mask * (-1) + 1) * output
+                    output = trick.finetune_mouth(source_img, target, output)
+                output = output.clamp(-1, 1)
+            elif 'simswap_official' in fs_model_name:
+                output = fs_model.image_infer(source_tensor=source_img, target_tensor=target)
+                output = output.clamp(-1, 1)
+            if isinstance(output, tuple):
+                target = output[0][0] * 0.5 + 0.5
+            else:
+                target = output[0] * 0.5 + 0.5
+        targets.append(np.array(tensor2pil_transform(target)))
+        Ms.append(M)
+        count += 1
+        if count > frames:
+            break
+    os.makedirs(out_path, exist_ok=True)
+    return targets, t_facial_masks, Ms, original_frames, names, fps
+
+
+def swap_image_gr(img1, img2, use_post=False, use_gpen=False, gpu_mode=True):
+    root_dir = make_abs_path("./online_data")
+    req_id = uuid.uuid1().hex
+    data_dir = os.path.join(root_dir, req_id)
+    os.makedirs(data_dir, exist_ok=True)
+    source_path = os.path.join(data_dir, "source.png")
+    target_path = os.path.join(data_dir, "target.png")
+    filename = "paste_back_out_target.png"
+    out_path = os.path.join(data_dir, filename)
+    cv2.imwrite(source_path, img1[:, :, ::-1])
+    cv2.imwrite(target_path, img2[:, :, ::-1])
+    swap_image(
+        source_path,
+        target_path,
+        data_dir,
+        T,
+        fs_model,
+        gpu_mode=gpu_mode,
+        align_target='ffhq',
+        align_source='ffhq',
+        use_post=use_post,
+        use_gpen=use_gpen,
+        in_size=in_size,
+    )
+    out = cv2.imread(out_path)[..., ::-1]
+    return out
+
+
+def swap_video_gr(img1, target_path, use_gpu=True, frames=9999999):
+    root_dir = make_abs_path("./online_data")
+    req_id = uuid.uuid1().hex
+    data_dir = os.path.join(root_dir, req_id)
+    os.makedirs(data_dir, exist_ok=True)
+    source_path = os.path.join(data_dir, "source.png")
+    cv2.imwrite(source_path, img1[:, :, ::-1])
+    out_dir = os.path.join(data_dir, "out")
+    out_name = "output.mp4"
+    targets, t_facial_masks, Ms, original_frames, names, fps = process_video(
+        source_path,
+        target_path,
+        out_dir,
+        T,
+        fs_model,
+        gpu_mode=use_gpu,
+        frames=frames,
+        align_target='ffhq',
+        align_source='ffhq',
+        use_tddfav2=False,
+    )
+
+    pool_process = 170
+    audio = True
+    concat = False
+
+    if pool_process <= 1:
+        for target, M, original_target, name, t_facial_mask in tqdm.tqdm(
+                zip(targets, Ms, original_frames, names, t_facial_masks)
+        ):
+            if M is None or target is None:
+                Image.fromarray(original_target.astype(np.uint8)).save(name)
+                continue
+            Image.fromarray(paste_back(np.array(target), M, original_target, t_facial_mask)).save(name)
+    else:
+        with Pool(pool_process) as pool:
+            pool.map(save, zip(targets, Ms, original_frames, names, t_facial_masks))
+
+    video_save_path = os.path.join(out_dir, out_name)
+    if audio:
+        print("use audio")
+        os.system(
+            f"ffmpeg  -y -r {fps} -i {out_dir}/frame_%05d.png -i {target_path}"
+            f" -map 0:v:0 -map 1:a:0? -c:a copy -c:v libx264 -r {fps} -crf 10 -pix_fmt yuv420p  {video_save_path}"
+        )
+    else:
+        print("no audio")
+        os.system(
+            f"ffmpeg  -y -r {fps} -i ./tmp/frame_%05d.png "
+            f"-c:v libx264 -r {fps} -crf 10 -pix_fmt yuv420p {video_save_path}"
+        )
+    # ffmpeg -i left.mp4 -i right.mp4 -filter_complex hstack output.mp4
+    if concat:
+        concat_video_save_path = os.path.join(out_dir, "concat_" + out_name)
+        os.system(
+            f"ffmpeg -y  -i {target_path}  -i {video_save_path} -filter_complex hstack {concat_video_save_path}"
+        )
+    # delete tmp file
+    shutil.rmtree("./tmp/")
+    for match in glob.glob(os.path.join(out_dir, "*.png")):
+        os.remove(match)
+    print(video_save_path)
+    return video_save_path
+
+
+if __name__ == "__main__":
+    with gr.Blocks() as demo:
+        gr.Markdown("SuperSwap")
+
+        with gr.Tab("Image"):
+            with gr.Row(equal_height=True):
+                with gr.Column(scale=3):
+                    image1_input = gr.Image()
+                    image2_input = gr.Image()
+                    use_post = gr.Checkbox(label="后处理")
+                    use_gpen = gr.Checkbox(label="超分增强")
+                with gr.Column(scale=2):
+                    image_output = gr.Image()
+                    image_button = gr.Button("换脸")
+        with gr.Tab("Video"):
+            with gr.Row(equal_height=True):
+                with gr.Column(scale=3):
+                    image3_input = gr.Image()
+                    video_input = gr.Video()
+                with gr.Column(scale=2):
+                    video_output = gr.Video()
+                    video_button = gr.Button("换脸")
+        image_button.click(
+            swap_image_gr,
+            inputs=[image1_input, image2_input, use_post, use_gpen],
+            outputs=image_output,
+        )
+        video_button.click(
+            swap_video_gr,
+            inputs=[image3_input, video_input],
+            outputs=video_output,
+        )
+
+    demo.launch(server_name="0.0.0.0", server_port=7860)