Update app.py

app.py

@@ -52,30 +52,30 @@ def text_image_to_image(image, prompt):
     modified_image.save(image_path)
     return modified_image, image_path
 
-def process_images(image_paths):
+def process_images(image_paths, remove_bg_method):
     with ThreadPoolExecutor() as executor:
-        results = list(executor.map(remove_background_rembg, image_paths))
+        if remove_bg_method == 'rembg':
+            results = list(executor.map(remove_background_rembg, image_paths))
+        elif remove_bg_method == 'bria':
+            results = list(executor.map(remove_background_bria, image_paths))
+        else:  # No background removal
+            results = image_paths  # Just return the original paths
     return results
 
 def get_bounding_box_with_threshold(image, threshold):
-    # Convert image to numpy array
     img_array = np.array(image)
+    alpha = img_array[:, :, 3]
     
-    # Get alpha channel
-    alpha = img_array[:,:,3]
-    
-    # Find rows and columns where alpha > threshold
     rows = np.any(alpha > threshold, axis=1)
     cols = np.any(alpha > threshold, axis=0)
     
-    # Find the bounding box
-    top, bottom = np.where(rows)[0][[0, -1]]
-    left, right = np.where(cols)[0][[0, -1]]
+    try:
+        top, bottom = np.where(rows)[0][[0, -1]]
+        left, right = np.where(cols)[0][[0, -1]]
+    except IndexError:
+        return None  # No non-transparent pixels found
     
-    if left < right and top < bottom:
-        return (left, top, right, bottom)
-    else:
-        return None
+    return (left, top, right, bottom) if left < right and top < bottom else None
 
 def check_cropped_sides(image, tolerance):
     cropped_sides = []
@@ -493,21 +493,19 @@ def create_canvas(canvas_size, bg_choice, custom_color):
     else:  # transparent
         return Image.new("RGBA", canvas_size, (0, 0, 0, 0))
 
-def apply_watermark(canvas, watermark_path):
+# Function to add watermark
+def add_watermark(image, watermark_path):
     watermark = Image.open(watermark_path).convert("RGBA")
-    
-    # Get dimensions of both the canvas and watermark
-    canvas_width, canvas_height = canvas.size
-    watermark_width, watermark_height = watermark.size
-
-    # Calculate the position to center the watermark
-    x = (canvas_width - watermark_width) // 2
-    y = (canvas_height - watermark_height) // 2
+    image = image.convert("RGBA")
 
-    # Paste the watermark on the canvas
-    canvas.paste(watermark, (x, y), watermark)
+    # Center the watermark
+    img_width, img_height = image.size
+    wm_width, wm_height = watermark.size
+    x = (img_width - wm_width) // 2
+    y = (img_height - wm_height) // 2
 
-    return canvas
+    image.paste(watermark, (x, y), watermark)
+    return image
     
 def save_image(canvas, output_folder, filename, output_format):
     output_ext = 'jpg' if output_format == 'JPG' else 'png'