import os

import gradio as gr
import numpy as np
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"
CHECKPOINT_FILENAME = "pifuhd.pt"
CHECKPOINT_PATH = f"checkpoints/{CHECKPOINT_FILENAME}"

pose_estimator = PoseEstimator("cpu")

if not os.path.exists(CHECKPOINT_PATH):
    checkpoint_path = hf_hub_download(repo_id=REPO_ID, filename=CHECKPOINT_FILENAME)
else:
    checkpoint_path = CHECKPOINT_PATH
cmd = [
    '--dataroot', './data',
    '--results_path', './results',
    '--loadSize', '1024',
    '--resolution', '512',
    '--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 predict(img: np.ndarray):
    bundle = make_bundle(img, "Model3D")
    dataset = EvalWMetaDataset(opts, [bundle])
    return reconstructor.evaluate(dataset)


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")
            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, output_img, output_obj], variant="stop")

    run_btn.click(predict, [input_img], [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()


import numpy as np
from human_pose_estimator.modules.pose import Pose


def rect(estimator, img: np.ndarray, height_size=512):
    num_keypoints = Pose.num_kpts
    _, pose_entries, all_keypoints = estimator.get_poses(img, height_size)

    rects = []
    for n in range(len(pose_entries)):
        if len(pose_entries[n]) == 0:
            continue
        pose_keypoints = np.ones((num_keypoints, 2), dtype=np.int32) * -1

        valid_keypoints = []
        for kpt_id in range(num_keypoints):
            if pose_entries[n][kpt_id] != -1.0:  # keypoint was found
                pose_keypoints[kpt_id, 0] = int(all_keypoints[int(pose_entries[n][kpt_id]), 0])
                pose_keypoints[kpt_id, 1] = int(all_keypoints[int(pose_entries[n][kpt_id]), 1])
                valid_keypoints.append([pose_keypoints[kpt_id, 0], pose_keypoints[kpt_id, 1]])
        valid_keypoints = np.array(valid_keypoints)

        if pose_entries[n][10] != -1.0 or pose_entries[n][13] != -1.0:
            pmin = valid_keypoints.min(0)
            pmax = valid_keypoints.max(0)

            center = (0.5 * (pmax[:2] + pmin[:2])).astype(np.int32)
            radius = int(0.65 * max(pmax[0] - pmin[0], pmax[1] - pmin[1]))
        elif pose_entries[n][10] == -1.0 and pose_entries[n][13] == -1.0 and pose_entries[n][8] != -1.0 and \
                pose_entries[n][11] != -1.0:
            # if leg is missing, use pelvis to get cropping
            center = (0.5 * (pose_keypoints[8] + pose_keypoints[11])).astype(np.int32)
            radius = int(1.45 * np.sqrt(((center[None, :] - valid_keypoints) ** 2).sum(1)).max(0))
            center[1] += int(0.05 * radius)
        else:
            center = np.array([img.shape[1] // 2, img.shape[0] // 2])
            radius = max(img.shape[1] // 2, img.shape[0] // 2)

        x1 = center[0] - radius
        y1 = center[1] - radius

        rects.append([x1, y1, 2 * radius, 2 * radius])

    return img, rects



У меня почему то зеркалит изображение и оно с оттенками синего цвета