sketch2pose / app.py
kbrodt's picture
Update app.py
be3490d
raw
history blame
2.11 kB
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)
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.glb"
if not mesh_path.is_file():
subprocess.call(cmd.split())
return str(mesh_path)
examples = []
for img_path in Path("./data/images").glob("*"):
examples.append([str(img_path), True, True, False, True])
break
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()