Update app.py

app.py

@@ -1,11 +1,12 @@
 import gradio as gr
-from PIL import Image
+from PIL import Image, ImageFilter
 import rembg
 import io
+import base64
 
 def process_image(image, outline_size):
     if image is None:
-        return None, None  # Return None for both image and file
+        return None
 
     # Convert image to bytes for rembg
     img_byte_arr = io.BytesIO()
@@ -16,35 +17,34 @@ def process_image(image, outline_size):
     processed_bytes = rembg.remove(img_byte_arr.getvalue())
     processed_img = Image.open(io.BytesIO(processed_bytes)).convert("RGBA")
 
-    # Add white outline
-    if outline_size > 0:
-        background = Image.new(
-            "RGBA",
-            (
-                processed_img.width + 2 * outline_size,
-                processed_img.height + 2 * outline_size,
-            ),
-            (0, 0, 0, 0),  # Transparent background
-        )
-        background.paste(processed_img, (outline_size, outline_size), processed_img)
-
-        # Draw white outline
-        for x in range(background.width):
-            for y in range(background.height):
-                if (x < outline_size or y < outline_size or x >= background.width - outline_size or y >= background.height - outline_size):
-                     if background.getpixel((x, y))[3] == 0:  #Check if pixel is originally transparent
-                        background.putpixel((x, y), (255, 255, 255, 255)) #Set to white if so
-        final_img = background
-
-    else:
-        final_img = processed_img
-
-    # Save as PNG to a BytesIO object
+    # Create outline mask
+    alpha = processed_img.split()[-1]
+    alpha = Image.merge('L', [alpha])
+    blurred = alpha.filter(ImageFilter.GaussianBlur(radius=outline_size))
+    new_alpha = blurred.point(lambda x: 0 if x == 0 else 255)  # Binary mask
+
+    # Apply outline to the transparent image
+    outlined_img = Image.new("RGBA", processed_img.size, (0, 0, 0, 0))  # Fully transparent
+    pixels = outlined_img.load()
+
+    alpha_pixels = new_alpha.load()
+    for x in range(outlined_img.width):
+        for y in range(outlined_img.height):
+            if alpha_pixels[x, y] > 0:
+                pixels[x, y] = (255, 255, 255, 255)  # White outline
+
+    # Overlay original image
+    outlined_img = Image.alpha_composite(outlined_img, processed_img)
+
+
+    # Save the result as bytes for Gradio
     output_buffer = io.BytesIO()
-    final_img.save(output_buffer, format="PNG")
+    outlined_img.save(output_buffer, format="PNG")
     output_buffer.seek(0)
+    img_data = base64.b64encode(output_buffer.read()).decode("utf-8")
+    data_url = f"data:image/png;base64,{img_data}"
 
-    return final_img, output_buffer
+    return data_url
 
 # Gradio Interface
 with gr.Blocks(title="Sticker Maker") as interface:
@@ -60,25 +60,38 @@ with gr.Blocks(title="Sticker Maker") as interface:
             create_btn = gr.Button("Create Sticker", variant="primary")
 
         with gr.Column():
-            output_image = gr.Image(label="Preview", type="pil", interactive=False)
-            download_link = gr.HTML("<a style='display: none;' id='download-link'>Download Sticker</a>")
-
-    def update_outputs(image, outline_size):
-        processed_image, image_bytes = process_image(image, outline_size)
-        if processed_image is None:
-            return None, "<a style='display: none;' id='download-link'>Download Sticker</a>"
+            # Use HTML to create an image and a download link
+            image_output_html = gr.HTML(
+                """
+                <div style="position: relative;">
+                    <img id="sticker-preview" src="" style="max-width: 100%; max-height: 400px;">
+                    <a id="download-link" style="position: absolute; top: 10px; right: 10px; background-color: rgba(255, 255, 255, 0.7); padding: 5px; border-radius: 5px; display:none;" download="sticker.png">Download</a>
+                </div>
+                """
+            )
 
-        # Create a data URL for the image
-        data_url = f"data:image/png;base64,{image_bytes.getvalue().hex()}"
+    def update_image(image, outline_size):
+        data_url = process_image(image, outline_size)
+        if data_url:
+            return f"""
+                <div style="position: relative;">
+                    <img id="sticker-preview" src="{data_url}" style="max-width: 100%; max-height: 400px;">
+                    <a id="download-link" href="{data_url}" download="sticker.png" style="position: absolute; top: 10px; right: 10px; background-color: rgba(255, 255, 255, 0.7); padding: 5px; border-radius: 5px;">Download</a>
+                </div>
+                """
+        else:
+            return """
+                <div style="position: relative;">
+                    <img id="sticker-preview" src="" style="max-width: 100%; max-height: 400px;">
+                    <a id="download-link" style="position: absolute; top: 10px; right: 10px; background-color: rgba(255, 255, 255, 0.7); padding: 5px; border-radius: 5px; display:none;" download="sticker.png">Download</a>
+                </div>
+                """
 
-        # Create the download link
-        download_html = f"<a href='{data_url}' download='sticker.png'>Download Sticker</a>"
-        return processed_image, download_html
 
     create_btn.click(
-        fn=update_outputs,
+        fn=update_image,
         inputs=[image_upload, outline_size],
-        outputs=[output_image, download_link],
+        outputs=image_output_html,
     )
 
 if __name__ == "__main__":