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((new_width, 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;
}
"""
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 height 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')
                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=643534)
                button = gr.Button()
            output_image = gr.Image(show_label=False, type="filepath", interactive=False, height=512, width=512)
      
    button.click(fn=generate, inputs=[image, textbox, seed], outputs=[output_image])

demo.queue(max_size=10).launch(inline=False, debug=True)