added white background

copaint/gradio_ui.py

@@ -24,6 +24,18 @@ logger = logging.getLogger(__name__)
 fromPIltoTensor = torchvision.transforms.ToTensor()
 fromTensortoPIL = torchvision.transforms.ToPILImage()
 
+from PIL import Image
+
+def remove_transparency(image: Image.Image):
+    # Convert transparency to white background
+    if image.mode in ("RGBA", "LA"):
+        background = Image.new("RGB", image.size, (255, 255, 255))  # white background
+        background.paste(image, mask=image.split()[-1])  # paste using alpha channel as mask
+        image = background
+    else:
+        image = image.convert("RGB")  # just to be safe
+
+    return image  # or continue processing
 
 def add_grid_to_image(image, h_cells, w_cells):
     # image is a torch tensor of shape (3, h, w)
@@ -177,8 +189,12 @@ def build_gradio_ui():
 
             # Upload Design Template
             with gr.Column(scale=2):
-                input_image = gr.Image(type="pil", label="Upload Your Design")
-
+                input_image = gr.Image(type="pil", label="Upload Your Design", preprocess=handle_upload)
+            image_input.upload(
+                fn=remove_transparency,
+                inputs=image_input,
+                outputs=image_input
+            )
             with gr.Column(scale=1):
                 # Grid
                 with gr.Tab("Grid Layout"):