Upload app.py

app.py

@@ -0,0 +1,96 @@
+import textwrap
+import subprocess
+import shutil
+import os
+from pathlib import Path
+
+import torch
+import gradio as gr
+from huggingface_hub import hf_hub_download
+
+
+REPO_ID = "kbrodt/sketch2pose"
+API_TOKEN = os.environ["sketch2pose"]
+ASSET_DIR = Path("./assets")
+
+filename = "models_smplx_v1_1.zip"
+smpl_path = hf_hub_download(
+    repo_id=REPO_ID,
+    repo_type="model",
+    filename=filename,
+    use_auth_token=API_TOKEN,
+    cache_dir=ASSET_DIR,
+)
+if not (ASSET_DIR / filename).is_file():
+    shutil.copy(smpl_path, ASSET_DIR)
+
+#filename = "output.zip"
+#output_path = hf_hub_download(
+    #repo_id=REPO_ID,
+    #repo_type="dataset",
+    #filename=filename,
+    #use_auth_token=API_TOKEN,
+    #cache_dir=ASSET_DIR,
+#)
+#if not (ASSET_DIR / filename).is_file():
+    #shutil.copy(output_path, ASSET_DIR)
+
+subprocess.run("bash ./scripts/download.sh".split())
+subprocess.run("bash ./scripts/prepare.sh".split())
+
+
+SAVE_DIR = "output"
+CMD = textwrap.dedent("""
+    python src/pose.py
+        --save-path {}
+        --img-path {}
+""")
+
+
+def main():
+    save_dir = Path(SAVE_DIR)
+    save_dir.mkdir(parents=True, exist_ok=True)
+
+    def pose(img_path, use_cos=True, use_angle_transf=True, use_contacts=False, use_natural=True):
+        if use_cos == False:
+            use_angle_transf = False
+
+        cmd = CMD.format(save_dir, img_path)
+        if use_cos:
+            cmd = cmd + " --use-cos"
+        if use_angle_transf:
+            cmd = cmd + " --use-angle-transf"
+        if use_contacts:
+            cmd = cmd + " --use-contacts"
+        if use_natural:
+            cmd = cmd + " --use-natural"
+
+        out_dir = (save_dir / Path(img_path).name).with_suffix("")
+        mesh_path = out_dir / "us.ply"
+
+        if not mesh_path.is_file():
+            subprocess.call(cmd.split())
+
+        return mesh_path
+
+    examples = []
+    for img_path in Path("./data/images").glob("*"):
+        examples.append([str(img_path), True, True, False, True])
+    demo = gr.Interface(
+        fn=pose,
+        inputs=[
+            gr.Image(type="filepath"),
+            gr.Checkbox(value=True),
+            gr.Checkbox(value=True),
+            gr.Checkbox(value=False, interactive=torch.cuda.is_available()),
+            gr.Checkbox(value=True),
+        ],
+        outputs=gr.Model3D(),
+        examples=examples,
+    )
+
+    demo.launch()
+
+
+if __name__ == "__main__":
+    main()