Update app.py

app.py

@@ -1,59 +1,19 @@
 import gradio as gr
-from PIL import Image, ImageOps
+from PIL import Image
 import rembg
 import io
-import numpy as np
+import os
+import shutil
 
-def remove_background_add_outline(image, outline_size=5, outline_color=(255, 255, 255)):
+def process_image(image, outline_size, outline_color, add_outline):
     if image is None:
         return None, None
     
-    # Convert image to bytes for rembg
-    img_byte_arr = io.BytesIO()
-    image.save(img_byte_arr, format='PNG')
-    img_byte_arr.seek(0)
-    
-    # Remove background
-    processed_bytes = rembg.remove(img_byte_arr.getvalue())
-    processed_img = Image.open(io.BytesIO(processed_bytes)).convert("RGBA")
-    
-    # Add outline
-    if outline_size > 0:
-        # Create a mask from the alpha channel
-        alpha = processed_img.split()[3]
-        
-        # Create a slightly larger mask by dilating
-        larger_mask = ImageOps.expand(alpha, outline_size, fill=255)
-        
-        # Create an outline image
-        outline_img = Image.new('RGBA', larger_mask.size, (0, 0, 0, 0))
-        outline_color_with_alpha = outline_color + (255,)  # Add alpha to color
-        
-        # Draw the outline
-        for x in range(larger_mask.width):
-            for y in range(larger_mask.height):
-                if larger_mask.getpixel((x, y)) == 255 and (x < outline_size or y < outline_size or 
-                                                          x >= processed_img.width or y >= processed_img.height or
-                                                          processed_img.getpixel((x-outline_size, y-outline_size))[3] == 0):
-                    outline_img.putpixel((x, y), outline_color_with_alpha)
-        
-        # Combine the outline with the original image
-        final_img = Image.new('RGBA', larger_mask.size, (0, 0, 0, 0))
-        final_img.paste(outline_img, (0, 0), outline_img)
-        final_img.paste(processed_img, (outline_size, outline_size), processed_img)
-    else:
-        final_img = processed_img
-    
-    # Prepare file for download
-    download_path = "sticker.png"
-    final_img.save(download_path, format="PNG")
-    
-    return final_img, download_path
-
-# Faster but simpler method for adding outline
-def add_outline_simple(image, outline_size=5, outline_color=(255, 255, 255)):
-    if image is None:
-        return None, None
+    # Convert hex color to RGB
+    r = int(outline_color[1:3], 16)
+    g = int(outline_color[3:5], 16)
+    b = int(outline_color[5:7], 16)
+    rgb_color = (r, g, b)
     
     # Convert image to bytes for rembg
     img_byte_arr = io.BytesIO()
@@ -64,54 +24,115 @@ def add_outline_simple(image, outline_size=5, outline_color=(255, 255, 255)):
     processed_bytes = rembg.remove(img_byte_arr.getvalue())
     processed_img = Image.open(io.BytesIO(processed_bytes)).convert("RGBA")
     
-    # Create a larger background image with the outline color
-    if outline_size > 0:
+    # Add outline if selected
+    if add_outline and outline_size > 0:
+        # Create a larger background image with the outline color
         background = Image.new('RGBA', 
                               (processed_img.width + 2*outline_size, 
                                processed_img.height + 2*outline_size),
-                              outline_color + (255,))
+                              (0, 0, 0, 0))  # Start with transparent background
+        
+        # Create outline mask
+        outline_img = Image.new('RGBA', background.size, rgb_color + (255,))
         
         # Paste the processed image onto the background
+        background.paste(outline_img, (0, 0), create_outline_mask(processed_img, outline_size))
         background.paste(processed_img, (outline_size, outline_size), processed_img)
         final_img = background
     else:
         final_img = processed_img
     
-    # Prepare file for download
-    download_path = "sticker.png"
-    final_img.save(download_path, format="PNG")
+    # Create output directory if it doesn't exist
+    os.makedirs("outputs", exist_ok=True)
+    
+    # Save as PNG with transparency preserved
+    output_path = "outputs/sticker.png"
+    final_img.save(output_path, format="PNG")
+    
+    # Create a copy for gradio to return (workaround for download issues)
+    download_path = "outputs/download_sticker.png"
+    shutil.copy2(output_path, download_path)
     
     return final_img, download_path
 
+def create_outline_mask(image, thickness):
+    """Create a mask for the outline based on the image alpha channel"""
+    # Get mask from alpha channel
+    mask = image.split()[3]
+    w, h = image.size
+    
+    expanded = Image.new('L', (w + 2*thickness, h + 2*thickness), 0)
+    expanded.paste(mask, (thickness, thickness))
+    
+    # Create outline mask
+    result = Image.new('L', expanded.size, 0)
+    for i in range(thickness):
+        # Expand the mask in all 8 directions for each thickness level
+        temp = Image.new('L', expanded.size, 0)
+        # Expand right
+        temp.paste(expanded, (-1, 0))
+        result = Image.composite(temp, result, temp)
+        # Expand left
+        temp = Image.new('L', expanded.size, 0)
+        temp.paste(expanded, (1, 0))
+        result = Image.composite(temp, result, temp)
+        # Expand up
+        temp = Image.new('L', expanded.size, 0)
+        temp.paste(expanded, (0, 1))
+        result = Image.composite(temp, result, temp)
+        # Expand down
+        temp = Image.new('L', expanded.size, 0)
+        temp.paste(expanded, (0, -1))
+        result = Image.composite(temp, result, temp)
+        # Expand diagonal directions
+        temp = Image.new('L', expanded.size, 0)
+        temp.paste(expanded, (-1, -1))
+        result = Image.composite(temp, result, temp)
+        temp = Image.new('L', expanded.size, 0)
+        temp.paste(expanded, (-1, 1))
+        result = Image.composite(temp, result, temp)
+        temp = Image.new('L', expanded.size, 0)
+        temp.paste(expanded, (1, -1))
+        result = Image.composite(temp, result, temp)
+        temp = Image.new('L', expanded.size, 0)
+        temp.paste(expanded, (1, 1))
+        result = Image.composite(temp, result, temp)
+    
+    # Subtract the original mask to get just the outline
+    for y in range(expanded.height):
+        for x in range(expanded.width):
+            if expanded.getpixel((x, y)) > 0:
+                result.putpixel((x, y), 0)
+                
+    return result
+
 # Gradio Interface
 with gr.Blocks(title="Sticker Maker") as interface:
     gr.Markdown("# Sticker Maker")
-    gr.Markdown("Remove background, add white outline, and download as transparent PNG")
+    gr.Markdown("Remove background, add optional outline, and download as transparent PNG")
     
     with gr.Row():
         with gr.Column():
             image_upload = gr.Image(label="Upload Image", type="pil")
+            add_outline = gr.Checkbox(label="Add Outline", value=True)
             outline_size = gr.Slider(label="Outline Thickness", minimum=0, maximum=20, value=5, step=1)
             outline_color = gr.ColorPicker(label="Outline Color", value="#FFFFFF")
-            create_btn = gr.Button("Create Sticker")
+            create_btn = gr.Button("Create Sticker", variant="primary")
         
         with gr.Column():
-            output_preview = gr.Image(label="Preview", type="pil")
+            output_preview = gr.Image(label="Preview (Transparent Background)", type="pil")
             output_file = gr.File(label="Download PNG")
     
-    def process_image(image, outline_size, outline_color):
-        # Convert hex color to RGB
-        r = int(outline_color[1:3], 16)
-        g = int(outline_color[3:5], 16)
-        b = int(outline_color[5:7], 16)
-        rgb_color = (r, g, b)
-        
-        # Use the simpler method for better performance
-        return add_outline_simple(image, int(outline_size), rgb_color)
+    # Make outline options visible only when add_outline is checked
+    add_outline.change(
+        fn=lambda x: gr.update(visible=x), 
+        inputs=add_outline, 
+        outputs=[outline_size, outline_color]
+    )
     
     create_btn.click(
         fn=process_image,
-        inputs=[image_upload, outline_size, outline_color],
+        inputs=[image_upload, outline_size, outline_color, add_outline],
         outputs=[output_preview, output_file]
     )