import os
import sys
import subprocess
import pkg_resources
def is_package_installed(package_name):
try:
pkg_resources.get_distribution(package_name)
return True
except pkg_resources.DistributionNotFound:
return False
if is_package_installed("mmcv"):
print("MMCV is installed.")
else:
print("MMCV is not installed. Build it from the source.")
os.environ["MMCV_WITH_OPS"] = "1"
os.environ["FORCE_MLU"] = "1"
subprocess.run(["pip", "install", "-e", "./mmcv"], check=True)
subprocess.run(["pip", "list"], check=True)
if is_package_installed("pytorch3d"):
print("pytorch3d is installed.")
else:
print("pytorch3d is not installed. Build it from the source.")
subprocess.run(["pip", "install", "-e", "./pytorch3d"], check=True)
if is_package_installed("MultiScaleDeformableAttention"):
print("MultiScaleDeformableAttention is installed.")
else:
print("MultiScaleDeformableAttention is not installed. Build it from the source.")
subprocess.run(["pip", "install", "-e", "./models/aios/ops"], check=True)
import os.path as osp
from pathlib import Path
import cv2
import gradio as gr
import torch
import math
import spaces
from huggingface_hub import hf_hub_download
hf_hub_download(repo_id="ttxskk/AiOS", filename="aios_checkpoint.pth", local_dir="/home/user/app/pretrained_models")
OUT_FOLDER = '/home/user/app/demo_out'
os.makedirs(OUT_FOLDER, exist_ok=True)
DEMO_CONFIG = '/home/user/app/config/aios_smplx_demo.py'
MODEL_PATH = '/home/user/app/pretrained_models/aios_checkpoint.pth'
@spaces.GPU(enable_queue=True, duration=300)
def infer(video_input, batch_size, threshold=0.5, num_person=1):
os.system(f'rm -rf {OUT_FOLDER}/*')
os.system(f'torchrun --nproc_per_node 1 \
main.py \
-c {DEMO_CONFIG} \
--options batch_size={batch_size} backbone="resnet50" num_person={num_person} threshold={threshold} \
--resume {MODEL_PATH} \
--eval \
--inference \
--inference_input {video_input} \
--to_vid \
--output_dir {OUT_FOLDER}')
video_path = os.path.join(OUT_FOLDER, 'demo_vid.mp4')
save_path_img = os.path.join(OUT_FOLDER, 'res_img')
save_path_mesh = os.path.join(OUT_FOLDER, 'mesh')
save_mesh_file = os.path.join(OUT_FOLDER, 'mesh.zip')
os.system(f'zip -r {save_mesh_file} {save_path_mesh}')
yield video_path, save_mesh_file
TITLE = """
AiOS: All-in-One-Stage Expressive Human Pose and Shape Estimation
Recover multiple expressive human pose and shape recovery from an RGB image without any additional requirements, such as an off-the-shelf detection model.
