Spaces:
Running
on
L40S
Running
on
L40S
File size: 5,671 Bytes
c6208fd 855986f 3c57e49 1fb2726 855986f 7bb79b2 d1cea68 5e42991 1fb2726 855986f 6ec3acd 79c74d8 855986f 1fb2726 855986f 6ec3acd 855986f ce18996 67aa81a 855986f d1cea68 855986f 67aa81a e1a1acd 67aa81a 855986f 59299a4 855986f 59299a4 855986f 67aa81a 855986f 44ec107 855986f 59299a4 855986f 44ec107 c5993c6 afa3edd 855986f 1fb2726 820cfa8 3d09012 820cfa8 3d09012 722bde9 3d09012 1fb2726 abbec01 820cfa8 3d09012 855986f 3d09012 855986f 44ec107 855986f 820cfa8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 |
import torch
import os
import shutil
import tempfile
import gradio as gr
from PIL import Image
from rembg import remove
import sys
import uuid
import subprocess
from glob import glob
import requests
from huggingface_hub import snapshot_download
# Download models
os.makedirs("ckpts", exist_ok=True)
snapshot_download(
repo_id = "pengHTYX/PSHuman_Unclip_768_6views",
local_dir = "./ckpts"
)
os.makedirs("smpl_related", exist_ok=True)
snapshot_download(
repo_id = "fffiloni/PSHuman-SMPL-related",
local_dir = "./smpl_related"
)
# Folder containing example images
examples_folder = "examples"
# Retrieve all file paths in the folder
images_examples = [
os.path.join(examples_folder, file)
for file in os.listdir(examples_folder)
if os.path.isfile(os.path.join(examples_folder, file))
]
def remove_background(input_url):
# Create a temporary folder for downloaded and processed images
temp_dir = tempfile.mkdtemp()
# Download the image from the URL
image_path = os.path.join(temp_dir, 'input_image.png')
try:
image = Image.open(input_url)
flipped_image = image.transpose(Image.FLIP_LEFT_RIGHT) # Mirror-flip the image
flipped_image.save(image_path)
except Exception as e:
shutil.rmtree(temp_dir)
return f"Error downloading or saving the image: {str(e)}"
# Run background removal
try:
unique_id = str(uuid.uuid4())
removed_bg_path = os.path.join(temp_dir, f'output_image_rmbg_{unique_id}.png')
img = Image.open(image_path)
result = remove(img)
result.save(removed_bg_path)
# Remove the input image to keep the temp directory clean
os.remove(image_path)
except Exception as e:
shutil.rmtree(temp_dir)
return f"Error removing background: {str(e)}"
return removed_bg_path, temp_dir
def run_inference(temp_dir, removed_bg_path):
# Define the inference configuration
inference_config = "configs/inference-768-6view.yaml"
pretrained_model = "./ckpts"
crop_size = 740
seed = 600
num_views = 7
save_mode = "rgb"
try:
# Run the inference command
subprocess.run(
[
"python", "inference.py",
"--config", inference_config,
f"pretrained_model_name_or_path={pretrained_model}",
f"validation_dataset.crop_size={crop_size}",
f"with_smpl=false",
f"validation_dataset.root_dir={temp_dir}",
f"seed={seed}",
f"num_views={num_views}",
f"save_mode={save_mode}"
],
check=True
)
# Retrieve the file name without the extension
removed_bg_file_name = os.path.splitext(os.path.basename(removed_bg_path))[0]
output_videos = glob(os.path.join(f"out/{removed_bg_file_name}", "*.mp4"))
return output_videos
except subprocess.CalledProcessError as e:
return f"Error during inference: {str(e)}"
def process_image(input_url):
# Remove background
result = remove_background(input_url)
if isinstance(result, str) and result.startswith("Error"):
raise gr.Error(f"{result}") # Return the error message if something went wrong
removed_bg_path, temp_dir = result # Unpack only if successful
# Run inference
output_video = run_inference(temp_dir, removed_bg_path)
if isinstance(output_video, str) and output_video.startswith("Error"):
shutil.rmtree(temp_dir)
raise gr.Error(f"{output_images}") # Return the error message if inference failed
shutil.rmtree(temp_dir) # Cleanup temporary folder
print(output_video)
return output_video[0]
def gradio_interface():
with gr.Blocks() as app:
gr.Markdown("# PSHuman: Photorealistic Single-image 3D Human Reconstruction using Cross-Scale Multiview Diffusion and Explicit Remeshing")
gr.HTML("""
<div style="display:flex;column-gap:4px;">
<a href="https://github.com/pengHTYX/PSHuman">
<img src='https://img.shields.io/badge/GitHub-Repo-blue'>
</a>
<a href="https://penghtyx.github.io/PSHuman/">
<img src='https://img.shields.io/badge/Project-Page-green'>
</a>
<a href="https://arxiv.org/pdf/2409.10141">
<img src='https://img.shields.io/badge/ArXiv-Paper-red'>
</a>
<a href="https://huggingface.co/spaces/fffiloni/PSHuman?duplicate=true">
<img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-sm.svg" alt="Duplicate this Space">
</a>
<a href="https://huggingface.co/fffiloni">
<img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/follow-me-on-HF-sm-dark.svg" alt="Follow me on HF">
</a>
</div>
""")
with gr.Row():
with gr.Column(scale=2):
input_image = gr.Image(
label="Image input",
type="filepath",
height=240
)
submit_button = gr.Button("Process")
gr.Examples(
examples = examples_folder,
inputs = [input_image],
examples_per_page = 4
)
output_video= gr.Video(label="Output Video", scale=3)
submit_button.click(process_image, inputs=[input_image], outputs=[output_video])
return app
# Launch the Gradio app
app = gradio_interface()
app.launch(show_api=False, show_error=True)
|