Spaces:
Running
Running
| 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() |