Spaces:
Sleeping
Sleeping
import os | |
import sys | |
import os.path as osp | |
from pathlib import Path | |
import cv2 | |
import gradio as gr | |
import torch | |
import math | |
import spaces | |
try: | |
import mmpose | |
except: | |
os.system('pip install /home/user/app/main/transformer_utils') | |
os.system('cp -rf /home/user/app/assets/conversions.py /home/user/.pyenv/versions/3.8.18/lib/python3.8/site-packages/torchgeometry/core/conversions.py') | |
DEFAULT_MODEL='smpler_x_h32' | |
OUT_FOLDER = '/home/user/app/demo_out' | |
os.makedirs(OUT_FOLDER, exist_ok=True) | |
num_gpus = 1 if torch.cuda.is_available() else -1 | |
print("!!!", torch.cuda.is_available()) | |
print(torch.cuda.device_count()) | |
print(torch.version.cuda) | |
index = torch.cuda.current_device() | |
print(index) | |
print(torch.cuda.get_device_name(index)) | |
from main.inference import Inferer | |
inferer = Inferer(DEFAULT_MODEL, num_gpus, OUT_FOLDER) | |
def infer(video_input, in_threshold=0.5, num_people="Single person", render_mesh=False): | |
os.system(f'rm -rf {OUT_FOLDER}/*') | |
multi_person = False if (num_people == "Single person") else True | |
cap = cv2.VideoCapture(video_input) | |
fps = math.ceil(cap.get(5)) | |
width = int(cap.get(3)) | |
height = int(cap.get(4)) | |
fourcc = cv2.VideoWriter_fourcc(*'mp4v') | |
video_path = osp.join(OUT_FOLDER, f'out.m4v') | |
final_video_path = osp.join(OUT_FOLDER, f'out.mp4') | |
video_output = cv2.VideoWriter(video_path, fourcc, fps, (width, height)) | |
success = 1 | |
frame = 0 | |
while success: | |
success, original_img = cap.read() | |
if not success: | |
break | |
frame += 1 | |
img, mesh_paths, smplx_paths = inferer.infer(original_img, in_threshold, frame, multi_person, not(render_mesh)) | |
video_output.write(img) | |
yield img, None, None, None | |
cap.release() | |
video_output.release() | |
cv2.destroyAllWindows() | |
os.system(f'ffmpeg -i {video_path} -c copy {final_video_path}') | |
#Compress mesh and smplx files | |
save_path_mesh = os.path.join(OUT_FOLDER, 'mesh') | |
save_mesh_file = os.path.join(OUT_FOLDER, 'mesh.zip') | |
os.makedirs(save_path_mesh, exist_ok= True) | |
save_path_smplx = os.path.join(OUT_FOLDER, 'smplx') | |
save_smplx_file = os.path.join(OUT_FOLDER, 'smplx.zip') | |
os.makedirs(save_path_smplx, exist_ok= True) | |
os.system(f'zip -r {save_mesh_file} {save_path_mesh}') | |
os.system(f'zip -r {save_smplx_file} {save_path_smplx}') | |
yield img, video_path, save_mesh_file, save_smplx_file | |
TITLE = '''<h1 align="center">SMPLer-X: Scaling Up Expressive Human Pose and Shape Estimation</h1>''' | |
VIDEO = ''' | |
<center><iframe width="960" height="540" | |
src="https://www.youtube.com/embed/DepTqbPpVzY?si=qSeQuX-bgm_rON7E"title="SMPLer-X: Scaling Up Expressive Human Pose and Shape Estimation" frameborder="0" allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture; web-share" allowfullscreen> | |
</iframe> | |
</center><br>''' | |
DESCRIPTION = ''' | |
<b>Official Gradio demo</b> for <a href="https://caizhongang.com/projects/SMPLer-X/"><b>SMPLer-X: Scaling Up Expressive Human Pose and Shape Estimation</b></a>.<br> | |
<p> | |
Note: You can drop a video at the panel (or select one of the examples) | |
to obtain the 3D parametric reconstructions of the detected humans. | |
</p> | |
''' | |
with gr.Blocks(title="SMPLer-X", css=".gradio-container") as demo: | |
gr.Markdown(TITLE) | |
gr.HTML(VIDEO) | |
gr.Markdown(DESCRIPTION) | |
with gr.Row(): | |
with gr.Column(): | |
video_input = gr.Video(label="Input video", elem_classes="video") | |
threshold = gr.Slider(0, 1.0, value=0.5, label='BBox detection threshold') | |
with gr.Column(scale=2): | |
num_people = gr.Radio( | |
choices=["Single person", "Multiple people"], | |
value="Single person", | |
label="Number of people", | |
info="Choose how many people are there in the video. Choose 'single person' for faster inference.", | |
interactive=True, | |
scale=1,) | |
gr.HTML("""<br/>""") | |
mesh_as_vertices = gr.Checkbox( | |
label="Render as mesh", | |
info="By default, the estimated SMPL-X parameters are rendered as vertices for faster visualization. Check this option if you want to visualize meshes instead.", | |
interactive=True, | |
scale=1,) | |
send_button = gr.Button("Infer") | |
gr.HTML("""<br/>""") | |
with gr.Row(): | |
with gr.Column(): | |
processed_frames = gr.Image(label="Last processed frame") | |
video_output = gr.Video(elem_classes="video") | |
with gr.Column(): | |
meshes_output = gr.File(label="3D meshes") | |
smplx_output = gr.File(label= "SMPL-X models") | |
# example_images = gr.Examples([]) | |
send_button.click(fn=infer, inputs=[video_input, threshold, num_people, mesh_as_vertices], outputs=[processed_frames, video_output, meshes_output, smplx_output]) | |
# with gr.Row(): | |
example_videos = gr.Examples([ | |
['/home/user/app/assets/01.mp4'], | |
['/home/user/app/assets/02.mp4'], | |
['/home/user/app/assets/03.mp4'], | |
['/home/user/app/assets/04.mp4'], | |
['/home/user/app/assets/05.mp4'], | |
['/home/user/app/assets/06.mp4'], | |
['/home/user/app/assets/07.mp4'], | |
['/home/user/app/assets/08.mp4'], | |
['/home/user/app/assets/09.mp4'], | |
], | |
inputs=[video_input, 0.5]) | |
#demo.queue() | |
demo.launch(debug=True) | |