import gradio as gr
from rembg import remove
import numpy as np
from PIL import Image
import uuid
import os

def remove_bg(input_image):
    try: print(input_image.name)
    except: pass

    if not os.path.exists("./temp"): os.mkdir("./temp")

    output_path = f'./temp/Removed_{uuid.uuid4()}.png'
    
    input_array = np.array(input_image)

    output_array = remove(input_array)

    output_img = Image.fromarray(output_array)

    output_img.save(output_path)
    return output_path

with gr.Blocks() as app:
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="Input Image")
            generate_btn = gr.Button(value="Remove Background")
        with gr.Column():
            output__image = gr.Image(label="Output Image")
    
    generate_btn.click(
        fn=remove_bg, 
        inputs=[input_image], 
        outputs=[output__image], 
        show_progress="full"
    )

app.launch()