import os
import gradio as gr
import numpy as np
from PIL import Image


def render(azimuth, elevation, theta, dist, category, unit):
    img_id = np.random.randint(0, 10000)

    os.system(f'python render.py --azimuth {azimuth} --elevation {elevation} --theta {theta} --dist {dist} --category {category} --unit {unit} --img_id {img_id}')

    img = Image.open(f'{img_id:05d}.png')
    os.system(f'rm {img_id:05d}.png')

    return np.array(img)

os.system('sh setup.sh')


with gr.Blocks() as demo:
    gr.Markdown('# Visualize object pose')
    gr.Markdown('This app runs on a free HuggingFace Space with no GPU support. Rendering an image generally takes a few seconds.')
    with gr.Row():
        with gr.Column(scale=1):
            azimuth_box = gr.Textbox(label="Azimuth", value="45")
            elevation_box = gr.Textbox(label="Elevation", value="15")
            theta_box = gr.Textbox(label="Theta", value="0")
            dist_box = gr.Textbox(label="Distance", value="4")
            category_radio = gr.Radio(["Aeroplane", "Bicycle", "Boat", "Bottle", "Bus", "Car", "Chair", "Diningtable", "Motorbike", "Sofa", "Train", "Tvmonitor"], value="Aeroplane")
            unit_radio = gr.Radio(["Degree", "Radian"], value="Degree")
            render_btn = gr.Button("Render")
        with gr.Column(scale=1):
            output = gr.Image(shape=(256, 256))
        render_btn.click(fn=render, inputs=[azimuth_box, elevation_box, theta_box, dist_box, category_radio, unit_radio], outputs=output)

demo.launch()