import gradio as gr
import subprocess
import os
from PIL import Image

def resize_image(image_path, target_height, output_path):
    # Open the image file
    with Image.open(image_path) as img:
        # Calculate the ratio to resize the image to the target height
        #ratio = target_height / float(img.size[1])
        # Calculate the new width based on the aspect ratio
        #new_width = int(float(img.size[0]) * ratio)
        # Resize the image
        resized_img = img.resize((512, target_height), Image.LANCZOS)
        # Save the resized image
        resized_img.save(output_path)
        return output_path


def generate(image, prompt, seed):
    print(image, prompt, seed)
    image_path = os.path.splitext(image)[0]
    image_name = os.path.basename(image_path)

    resized=resize_image(image, 512, f"output/{image_name}.jpg")
    print(f"IMAGE NAME: {image_name}")
    command = f"python handrefiner.py --input_img {resized} --out_dir output --strength 0.55 --weights models/inpaint_depth_control.ckpt --prompt '{prompt}' --seed {seed}"
    try:
        result = subprocess.run(command, shell=True, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        output_path = 'output'
        print("Output:", result.stdout)
        print(output_path)
        # List all files and directories in the given directory
        contents = os.listdir("output")
        
        # Print the contents
        for item in contents:
            print(item)
        
        return f"output/{image_name}_0.jpg"
    except subprocess.CalledProcessError as e:
        print("Error:", e.stderr)
        return None

css="""
#col-container{
    max-width: 860px;
    margin: 0 auto;
}
#project-links{
    margin: 0 0 12px !important;
    column-gap: 8px;
    display: flex;
    justify-content: center;
    flex-wrap: nowrap;
    flex-direction: row;
    align-items: center;
}
"""
with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        
        gr.HTML("""
        <h2 style="text-align: center;">
            HandRefiner
        </h2>
        <p style="text-align: center;">
            Refining Malformed Hands in Generated Images by Diffusion-based Conditional Inpainting <br />
            For demo purpose, every input images are resized to 512 1:1 aspect ratio
        </p>
        <p style="margin:12px auto;display: flex;justify-content: center;">
            <a href="https://huggingface.co/spaces/fffiloni/HandRefiner?duplicate=true"><img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-lg.svg" alt="Duplicate this Space"></a>
        </p>
        """)
        
        with gr.Row():
            with gr.Column():
                image = gr.Image(type='filepath', sources=["upload"])
                textbox = gr.Textbox(show_label=False, value="a person facing the camera, making a hand gesture, indoor")
                seed = gr.Slider(label="Seed", minimum=0, maximum=1000000, value=42, step=1)
                submit_btn = gr.Button("Submit")
                gr.Examples(
                    examples = [
                        "examples/IMG_1050.jpeg",
                        "examples/IMG_1051.jpeg",
                        "examples/IMG_1052.jpeg",
                        "examples/IMG_1053.jpeg"
                    ],
                    inputs = [image]
                )
            with gr.Column():
                output_image = gr.Image(label="Fixed hands result")
                gr.HTML("""
                <p id="project-links" align="center">
                  <a href='https://github.com/wenquanlu/HandRefiner'><img src='https://img.shields.io/badge/Project-Page-Green'></a> <a href='https://arxiv.org/abs/2311.17957'><img src='https://img.shields.io/badge/Paper-Arxiv-red'></a>
                </p>
                <img src="https://github.com/wenquanlu/HandRefiner/raw/main/Figs/banner.png" style="margin: 0 auto;border-radius: 10px;" />
                """)
      
    submit_btn.click(fn=generate, inputs=[image, textbox, seed], outputs=[output_image])

demo.queue(max_size=10).launch(debug=True)