import os

os.system("pip install gradio==3.50")

import gradio as gr
from PIL import Image
import subprocess

# Install necessary packages
os.system("pip install 'git+https://github.com/facebookresearch/detectron2.git'")

# Clone the repository
os.system("git clone https://github.com/ShuhongChen/bizarre-pose-estimator.git")
os.chdir("bizarre-pose-estimator")

# Download necessary files
os.system("wget https://i.imgur.com/IkJzlaE.jpeg")
os.system("gdown https://drive.google.com/uc?id=17N5PutpYJTlKuNB6bdDaiQsPSIkYtiPm")

# Unzip and move the model files
os.system("unzip bizarre_pose_models.zip")
os.system("cp -a ./bizarre_pose_models/. .")

def inference(img):
    # Save the input image
    img.save("_input.png")

    # Run the pose estimator
    os.system("python3 -m _scripts.pose_estimator _input.png ./_train/character_pose_estim/runs/feat_concat+data.ckpt")

    # Load and return the output image
    return Image.open("./_samples/character_pose_estim.png")

title = "bizarre-pose-estimator"
description = "Gradio demo for Transfer Learning for Pose Estimation of Illustrated Characters. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."

article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2108.01819' target='_blank'>Transfer Learning for Pose Estimation of Illustrated Characters</a> | <a href='https://github.com/ShuhongChen/bizarre-pose-estimator' target='_blank'>Github Repo</a></p>"

examples=[["IkJzlaE.jpeg"]]
gr.Interface(
    inference, 
    gr.inputs.Image(type="file", label="Input"), 
    gr.outputs.Image(type="file", label="Output"),
    title=title,
    description=description,
    article=article,
    allow_flagging="never",
    examples=examples,
    enable_queue=True
).launch()