from rembg import remove
from PIL import Image, ImageOps
import gradio as gr


def change_color(passport='Passport Photo', color='Background Color', background='Photo Background'):
    if background:
        image = remove(passport)
    else:
        image = passport
    new_image = Image.new("RGB", image.size, color)
    new_image.paste(image, mask=image)
    # Add a black border of 10 pixels on all sides
    new_image = ImageOps.expand(new_image, border=5, fill='black')
    new_image = ImageOps.expand(new_image, border=1, fill='white')

    # Get the image size
    width, height = new_image.size

    # Create a new image to hold the photos
    output_image = Image.new("RGB", (width * 4, height * 2))

    # Paste the source image into the new image
    for row in range(2):
        for col in range(4):
            output_image.paste(new_image, (col * width, row * height))
    # Save the output image
    return output_image


inputs = [
  gr.Image(label="Passport Photo", type="pil", image_mode="RGBA"),
  gr.ColorPicker(label="Background Color",value="#ff00ff"),
  gr.Checkbox(label="Select if you want to remove background", value=True)
]
outputs = gr.Image(label="Download Image")

demo = gr.Interface(
  fn=change_color,
  inputs=inputs,
  outputs=outputs,
)

if __name__ == "__main__":
  demo.launch()