add round overlay

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import numpy as np
 import cv2
-from PIL import Image, ImageOps
+from PIL import Image, ImageOps, ImageDraw
 import os
 import torch
 from transformers import AutoModelForImageSegmentation
@@ -11,7 +11,6 @@ import re
 import urllib.request as urllib2
 from loguru import logger
 
-
 # Set up model and transformations
 def get_background_removal_model():
     try:
@@ -27,7 +26,6 @@ def get_background_removal_model():
         print(f"Error loading background removal model: {e}")
         return None, None
 
-
 # Set up image transformation
 transform_image = transforms.Compose(
     [
@@ -40,7 +38,6 @@ transform_image = transforms.Compose(
 # Cache for storing background removal results
 bg_removal_cache = {}
 
-
 def get_image_hash(image):
     """Generate a hash for an image to use as cache key"""
     if image is None:
@@ -53,7 +50,6 @@ def get_image_hash(image):
     # Include image dimensions in the hash to ensure uniqueness
     return f"{img_hash}_{image.width}_{image.height}"
 
-
 def remove_background(image, model_data):
     if model_data[0] is None:
         return None, None
@@ -103,7 +99,6 @@ def remove_background(image, model_data):
         logger.error(f"Error during background removal: {e}")
         return None, None
 
-
 def parse_color(color_str):
     """Parse different color formats including rgba strings"""
     if isinstance(color_str, tuple):
@@ -145,7 +140,6 @@ def parse_color(color_str):
     # Default fallback
     return (255, 255, 255, 255)  # White
 
-
 def add_person_border(image, mask, border_size, border_color="white"):
     """Add a border around the person based on the segmentation mask"""
     if border_size == 0:
@@ -177,7 +171,6 @@ def add_person_border(image, mask, border_size, border_color="white"):
 
     return result
 
-
 def detect_face(image):
     """Detect the largest face in the image and return its bounding box"""
     logger.info("Starting face detection")
@@ -211,7 +204,6 @@ def detect_face(image):
     logger.info(f"Largest face detected at: {largest_face}")
     return largest_face
 
-
 def center_portrait(portrait, face_box, target_width, target_height, zoom_level=1.0):
     """Center the portrait based on face position and crop to avoid blurriness"""
     if face_box is None:
@@ -249,9 +241,8 @@ def center_portrait(portrait, face_box, target_width, target_height, zoom_level=
 
     return centered_img, (offset_x, offset_y)
 
-
 def process_portrait(
-    input_image, border_size=10, bg_color="#0000FF", zoom_level=1.0, erode_size=5
+    input_image, border_size=10, bg_color="#0000FF", zoom_level=1.0, erode_size=5, circular_overlay=False
 ):
     if input_image is None:
         return None
@@ -327,12 +318,20 @@ def process_portrait(
     bottom = top + square_size
     final_image = final_image.crop((left, top, right, bottom))
 
+    if circular_overlay:
+        # Create a circular mask
+        mask = Image.new("L", (square_size, square_size), 0)
+        draw = ImageDraw.Draw(mask)
+        draw.ellipse((0, 0, square_size, square_size), fill=255)
+
+        # Apply the circular mask to the final image
+        final_image.putalpha(mask)
+
     logger.info(
         f"Processing complete (portrait offset by {offset}, zoom: {zoom_level})"
     )
     return final_image
 
-
 # Create Gradio interface
 with gr.Blocks(title="Cool Avatar Creator") as app:
     gr.Markdown("# Cool Avatar Creator")
@@ -353,6 +352,7 @@ with gr.Blocks(title="Cool Avatar Creator") as app:
             erode_slider = gr.Slider(
                 minimum=1, maximum=30, value=15, step=1, label="Erode Size"
             )
+            circular_overlay_toggle = gr.Checkbox(label="Enable Circular Overlay")
             process_button = gr.Button("Process Image")
 
         with gr.Column():
@@ -384,7 +384,7 @@ with gr.Blocks(title="Cool Avatar Creator") as app:
 
     process_button.click(
         fn=process_portrait,
-        inputs=[input_image, border_slider, bg_color, zoom_slider, erode_slider],
+        inputs=[input_image, border_slider, bg_color, zoom_slider, erode_slider, circular_overlay_toggle],
         outputs=output_image,
     )