pifuhd / app.py
punithkumar779's picture
h
b3eb10d verified
raw
history blame
3.25 kB
import gradio as gr
import numpy as np
import cv2
from PIL import Image
from PIFuHD.data import EvalWMetaDataset
from PIFuHD.data.ImageBundle import ImageBundle
from PIFuHD.options import BaseOptions
from PIFuHD.recontructor import Reconstructor
from huggingface_hub import hf_hub_download
from human_pose_estimator import PoseEstimator
from estimator import rect
REPO_ID = "cxeep/PIFuHD"
pose_estimator = PoseEstimator("cpu")
checkpoint_path = hf_hub_download(repo_id=REPO_ID, filename="pifuhd.pt")
cmd = [
'--dataroot', './data',
'--results_path', './results',
'--loadSize', '1024',
'--resolution', '256',
'--load_netMR_checkpoint_path', checkpoint_path,
'--start_id', '-1',
'--end_id', '-1'
]
parser = BaseOptions()
opts = parser.parse(cmd)
reconstructor = Reconstructor(opts)
def make_bundle(image, name):
image, rects = rect(pose_estimator, image)
return ImageBundle(img=image, name=name, meta=rects)
def process_video(video_path):
frames = []
cap = cv2.VideoCapture(video_path)
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
frames.append(Image.fromarray(frame))
cap.release()
models = []
for frame in frames:
bundle = make_bundle(np.array(frame), "Model3D")
dataset = EvalWMetaDataset(opts, [bundle])
model = reconstructor.evaluate(dataset)
models.append(model)
# TODO: Combine models into animation
output_animation = models[0] # Placeholder, replace with actual animation
return output_animation
def predict(input_img, input_video):
if input_video:
return process_video(input_video)
elif input_img is not None:
bundle = make_bundle(input_img, "Model3D")
dataset = EvalWMetaDataset(opts, [bundle])
return reconstructor.evaluate(dataset)
else:
return None, None
footer = r"""
<center>
<b>Demo for <a href='https://github.com/facebookresearch/pifuhd'>PIFuHD</a></b>
</center>
"""
with gr.Blocks(title="PIFuHD") as app:
gr.HTML("<center><h1>3D Human Digitization</h1></center>")
gr.HTML("<center><h3>PIFuHD: Multi-Level Pixel-Aligned Implicit Function for High-Resolution 3D Human Digitization (CVPR 2020)</h3></center>")
with gr.Row(equal_height=False):
with gr.Column():
input_img = gr.Image(type="numpy", label="Input image")
input_video = gr.Video(label="Input Video")
run_btn = gr.Button(variant="primary")
with gr.Column():
output_obj = gr.Model3D(label="Output model")
output_img = gr.Image(type="filepath", label="Output image")
gr.ClearButton(components=[input_img, input_video, output_img, output_obj], variant="stop")
run_btn.click(predict, [input_img, input_video], [output_img, output_obj])
with gr.Row():
blobs = [[f"examples/{x:02d}.png"] for x in range(1, 4)]
examples = gr.Dataset(components=[input_img], samples=blobs)
examples.click(lambda x: x[0], [examples], [input_img])
with gr.Row():
gr.HTML(footer)
app.launch(share=False, debug=True, show_error=True)
app.queue()