import os
try:
    os.system("pip install --upgrade  torch==1.11.0+cu113 torchvision==0.12.0+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html")
except Exception as e:
    print(e)

from pydoc import describe
from huggingface_hub import hf_hub_download
import gradio as gr
import subprocess
import os
import datetime
from PIL import Image
from remove_bg import RemoveBackground
import torch, torchvision, skimage 

print(
    "torch: ", torch.__version__,
    "\ntorchvision: ",torchvision.__version__,
    "\nskimage:", skimage.__version__
)

net_C = hf_hub_download("radames/PIFu-upright-standing", filename="net_C")
net_G = hf_hub_download("radames/PIFu-upright-standing", filename="net_G")
torch.hub.load('pytorch/vision:v0.10.0', 'deeplabv3_resnet101', pretrained=True)

remove_bg = RemoveBackground()

env = {
    **os.environ,
    "CHECKPOINTS_NETG_PATH": net_G,
    "CHECKPOINTS_NETC_PATH": net_C,
    "RESULTS_PATH": './results',
}


def process(img_path):
    base = os.path.basename(img_path)
    img_name = os.path.splitext(base)[0]
    print("image name", img_name)
    img_raw = Image.open(img_path)
    img = img_raw.resize(
        (800, int(800 * img_raw.size[1] / img_raw.size[0])),
        Image.Resampling.LANCZOS)

    # remove background
    print("Removeing background")
    try: 
        foreground = Image.fromarray(remove_bg.inference(img), 'RGBA')
        foreground.save("./PIFu/inputs/" + img_name + ".png")
    except Exception as e:
        print(e)
    print("Aliging mask with input training image")
    subprocess.Popen(["python", "./apps/crop_img.py", "--input_image",
                      f'./inputs/{img_name}.png', "--out_path", "./inputs"], cwd="PIFu").communicate()

    print("Generating 3D model")
    subprocess.Popen("./scripts/test.sh", env={
                     **env,
                     "INPUT_IMAGE_PATH": f'./inputs/{img_name}.png',
                     "VOL_RES": "128"},
                     cwd="PIFu").communicate()
    print("DONE 3D model")
    return f'./PIFu/results/spaces_demo/result_{img_name}.glb'


examples = [["./examples/" + img] for img in sorted(os.listdir("./examples/"))]
description = '''
# PIFu Clothed Human Digitization
#### PIFu: Pixel-Aligned Implicit Function for High-Resolution Clothed Human Digitization
<base target="_blank">

This is a demo for <a href="https://github.com/shunsukesaito/PIFu" target="_blank"> PIFu model </a>.
The pre-trained model has the following warning:
> Warning: The released model is trained with mostly upright standing scans with weak perspectie projection and the pitch angle of 0 degree. Reconstruction quality may degrade for images highly deviated from trainining data.

**The inference takes about 180seconds for a new image.**

<details> 
<summary>More</summary>

#### Image Credits

* Julien and Clem
* [StyleGAN Humans](https://huggingface.co/spaces/hysts/StyleGAN-Human)
* [Renderpeople: Dennis](https://renderpeople.com)


#### More
* https://phorhum.github.io/
* https://github.com/yuliangxiu/icon
* https://shunsukesaito.github.io/PIFuHD/

</details>
'''

iface = gr.Interface(
    fn=process,
    description=description,
    inputs=gr.Image(type="filepath", label="Input Image"),
    outputs=gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0],  label="3D Model"),
    examples=examples,
    allow_flagging="never",
    cache_examples=True
)

if __name__ == "__main__":
    iface.launch(debug=True, enable_queue=False)