main.py

# Dependencies
"""
1) gradio==3.4
2) httpx==0.32.2
3) setuptools
"""
import gradio as gr
import imageio

def Mask(img):
    """
    Function to process the input image and generate a mask.
    
    Args:
    img (dict): Dictionary containing the base image and the mask image.
    
    Returns:
    tuple: A tuple containing the base image and the mask image.
    """
    try:
        # Save the mask image to a file
        imageio.imwrite("output_image.png", img["mask"])
        return img["image"], img["mask"]
    except KeyError as e:
        # Handle case where expected keys are not in the input dictionary
        return f"Key error: {e}", None
    except Exception as e:
        # Handle any other unexpected errors
        return f"An error occurred: {e}", None
def main():
    # Create the Gradio interface
    with gr.Blocks() as demo:
        with gr.Row():
            img = gr.Image(tool="sketch", label="Base Image", show_label=True)
            img1 = gr.Image()
            img2 = gr.Image(label="Mask Image", show_label=True)
        btn = gr.Button(label="Generate Mask")

        # Set the button click action
        btn.click(Mask, inputs=img, outputs=[img1, img2])

    # Launch the Gradio interface
    demo.launch(debug=True,share=True,server_port=8888)


if __name__=='__main__':
    main()