inital commit

.gitmodules

@@ -0,0 +1,9 @@
+[submodule "ROME"]
+    path = ROME
+	url = https://github.com/khakhulin/DECA/tree/1cc2361a2929a206e1b0330ee8b89fcda478d037
+[submodule "DECA"]
+	path = DECA
+	url = https://github.com/khakhulin/DECA/tree/1cc2361a2929a206e1b0330ee8b89fcda478d037
+[submodule "MODNet"]
+	path = MODNet
+	url = https://github.com/ZHKKKe/MODNet/tree/28165a451e4610c9d77cfdf925a94610bb2810fb

app.py

@@ -0,0 +1,51 @@
+import os, sys
+import numpy as np
+import torch
+from torchvision import transforms
+import matplotlib.pyplot as plt
+from PIL import Image
+
+# loading models ---- create model repo
+from huggingface_hub import hf_hub_url
+default_modnet_path = hf_hub_url('Pie31415/rome','modnet_photographic_portrait_matting.ckpt')
+default_model_path = hf_hub_url('Pie31415/rome','models/rome.pth')
+
+# parser configurations
+parser = argparse.ArgumentParser(conflict_handler='resolve')
+parser.add_argument('--save_dir', default='.', type=str)
+parser.add_argument('--save_render', default='True', type=args_utils.str2bool, choices=[True, False])
+parser.add_argument('--model_checkpoint', default=default_model_path, type=str)
+parser.add_argument('--modnet_path', default=default_modnet_path, type=str)
+parser.add_argument('--random_seed', default=0, type=int)
+parser.add_argument('--debug', action='store_true')
+parser.add_argument('--verbose', default='False', type=args_utils.str2bool, choices=[True, False])
+args, _ = parser.parse_known_args()
+
+parser = importlib.import_module(f'src.rome').ROME.add_argparse_args(parser)
+args = parser.parse_args()
+args.deca_path = 'DECA'
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+from infer import Infer
+
+infer = Infer(args)
+infer = infer.to(device)
+
+def predict(source_img, driver_img):
+    out = infer.evaluate(source_img, driver_img, crop_center=False)
+    res = tensor2image(torch.cat([out['source_information']['data_dict']['source_img'][0].cpu(),
+                              out['source_information']['data_dict']['target_img'][0].cpu(),
+                        out['render_masked'].cpu(), out['pred_target_shape_img'][0].cpu()], dim=2))
+    return res[..., ::-1]
+
+
+import gradio as gr
+gr.Interface(
+    fn=predict,
+    inputs=[
+        gr.Image(type="pil"),
+        gr.Image(type="pil")
+    ],
+    outputs=gr.Image(),
+    examples=[]).launch()