add bria for remove background model

.gitignore

@@ -0,0 +1 @@
+venv/

app.py

@@ -6,8 +6,10 @@ import io
 from rembg import remove
 import gradio as gr
 from concurrent.futures import ThreadPoolExecutor
+from transformers import pipeline
 
 def resize_and_crop_image(image_path, target_size=(1080, 1080), crop_mode='center'):
+    print(f"Resizing and cropping image: {image_path}")
     with Image.open(image_path) as img:
         width, height = img.size
 
@@ -42,69 +44,81 @@ def resize_and_crop_image(image_path, target_size=(1080, 1080), crop_mode='cente
 
         return cropped_img
 
-def remove_background(input_path):
+def remove_background_rembg(input_path):
+    print(f"Removing background using rembg: {input_path}")
     with open(input_path, 'rb') as i:
         input_image = i.read()
     output_image = remove(input_image)
     img = Image.open(io.BytesIO(output_image)).convert("RGBA")
     return img
 
-def process_single_image(image_path, output_folder, crop_mode, remove_bg, output_format, bg_choice, watermark_path=None):
+def remove_background_bria(input_path):
+    print(f"Removing background using bria: {input_path}")
+    pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True)
+    pillow_image = pipe(input_path)  # applies mask on input and returns a pillow image
+    return pillow_image
+
+def process_single_image(image_path, output_folder, crop_mode, remove_bg, bg_method, output_format, bg_choice, watermark_path=None):
     filename = os.path.basename(image_path)
     try:
+        print(f"Processing image: {filename}")
         if remove_bg == 'yes':
-            # Remove background
-            image_with_no_bg = remove_background(image_path)
+            if bg_method == 'rembg':
+                # Remove background using rembg
+                image_with_no_bg = remove_background_rembg(image_path)
+            elif bg_method == 'bria':
+                # Remove background using bria
+                image_with_no_bg = remove_background_bria(image_path)
+                
             temp_image_path = os.path.join(output_folder, f"temp_{filename}")
             image_with_no_bg.save(temp_image_path, format='PNG')
+            print(f"Background removed and saved temporary image: {temp_image_path}")
         else:
             temp_image_path = image_path
 
         # Resize and crop the image with or without background removal
         new_image = resize_and_crop_image(temp_image_path, crop_mode=crop_mode)
+        print(f"Resized and cropped image: {temp_image_path}")
+
+        # Save both versions of the image (with and without watermark)
+        images_paths = []
+
+        # Save without watermark
+        output_ext = 'jpg' if output_format == 'JPG' else 'png'
+        output_path_without_watermark = os.path.join(output_folder, f"without_watermark_{os.path.splitext(filename)[0]}.{output_ext}")
+        if output_format == 'JPG':
+            new_image.convert('RGB').save(output_path_without_watermark, format='JPEG')
+        else:
+            new_image.save(output_path_without_watermark, format='PNG')
+        images_paths.append(output_path_without_watermark)
+        print(f"Saved final image without watermark: {output_path_without_watermark}")
 
-        # Apply watermark if provided
+        # Apply watermark if provided and save the version with watermark
         if watermark_path:
             watermark = Image.open(watermark_path).convert("RGBA")
-            new_image.paste(watermark, (0, 0), watermark)
-
-        # Save the final image
-        if remove_bg == 'yes' and output_format == 'PNG' and bg_choice == 'transparent':
-            output_path = os.path.join(output_folder, f"{os.path.splitext(filename)[0]}.png")
-            new_image.save(output_path, format='PNG')
-        elif remove_bg == 'yes' and output_format == 'PNG' and bg_choice == 'white':
-            # Create a white background
-            background = Image.new("RGBA", new_image.size, (255, 255, 255, 255))
-            # Paste the image with no background onto the white background
-            background.paste(new_image, mask=new_image.split()[3])  # 3 is the alpha channel
-            output_path = os.path.join(output_folder, f"{os.path.splitext(filename)[0]}.png")
-            background.convert('RGB').save(output_path, format='PNG')
-        elif remove_bg == 'yes' and output_format == 'JPG':
-            # Create a white background
-            background = Image.new("RGB", new_image.size, (255, 255, 255))
-            # Paste the image with no background onto the white background
-            background.paste(new_image, mask=new_image.split()[3])  # 3 is the alpha channel
-            output_path = os.path.join(output_folder, f"{os.path.splitext(filename)[0]}.jpg")
-            background.save(output_path, format='JPEG')
-        else:
-            output_ext = 'jpg' if output_format == 'JPG' else 'png'
-            output_path = os.path.join(output_folder, f"{os.path.splitext(filename)[0]}.{output_ext}")
+            new_image_with_watermark = new_image.copy()
+            new_image_with_watermark.paste(watermark, (0, 0), watermark)
+            output_path_with_watermark = os.path.join(output_folder, f"with_watermark_{os.path.splitext(filename)[0]}.{output_ext}")
             if output_format == 'JPG':
-                new_image.convert('RGB').save(output_path, format='JPEG')
+                new_image_with_watermark.convert('RGB').save(output_path_with_watermark, format='JPEG')
             else:
-                new_image.save(output_path, format='PNG')
+                new_image_with_watermark.save(output_path_with_watermark, format='PNG')
+            images_paths.append(output_path_with_watermark)
+            print(f"Saved final image with watermark: {output_path_with_watermark}")
 
         if remove_bg == 'yes':
             # Remove the temporary file
             os.remove(temp_image_path)
+            print(f"Removed temporary file: {temp_image_path}")
 
-        return output_path
+        return images_paths
 
     except Exception as e:
         print(f"Error processing {filename}: {e}")
         return None
 
-def process_images(zip_file, crop_mode='center', remove_bg='yes', watermark_path=None, output_format='PNG', bg_choice='transparent', progress=gr.Progress()):
+def process_images(zip_file, crop_mode='center', remove_bg='yes', bg_method='rembg', watermark_path=None, output_format='PNG', bg_choice='transparent', progress=gr.Progress()):
+    print("Starting to process images...")
     # Create a temporary directory
     input_folder = "temp_input"
     output_folder = "temp_output"
@@ -116,42 +130,55 @@ def process_images(zip_file, crop_mode='center', remove_bg='yes', watermark_path
     os.makedirs(output_folder)
     
     # Extract the zip file
+    print(f"Extracting zip file: {zip_file}")
     with zipfile.ZipFile(zip_file, 'r') as zip_ref:
         zip_ref.extractall(input_folder)
     
     processed_images = []
     image_files = [os.path.join(input_folder, f) for f in os.listdir(input_folder) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif'))]
     total_images = len(image_files)
+    print(f"Found {total_images} images to process.")
 
     # Process images using ThreadPoolExecutor
     with ThreadPoolExecutor(max_workers=2) as executor:
-        future_to_image = {executor.submit(process_single_image, image_path, output_folder, crop_mode, remove_bg, output_format, bg_choice, watermark_path): image_path for image_path in image_files}
+        future_to_image = {executor.submit(process_single_image, image_path, output_folder, crop_mode, remove_bg, bg_method, output_format, bg_choice, watermark_path): image_path for image_path in image_files}
         for idx, future in enumerate(future_to_image):
             result = future.result()
             if result:
-                processed_images.append(result)
+                processed_images.extend(result)
             # Update progress
             progress((idx + 1) / total_images, f"{idx + 1}/{total_images} images processed")
-    
+            print(f"Processed {idx + 1}/{total_images} images")
+
     # Create a zip file of the processed images
     output_zip_path = "processed_images.zip"
     with zipfile.ZipFile(output_zip_path, 'w') as zipf:
         for file in processed_images:
-            zipf.write(file, os.path.basename(file))
-    
+            if "with_watermark" in file:
+                zipf.write(file, os.path.join("with_watermark", os.path.basename(file)))
+            else:
+                zipf.write(file, os.path.join("without_watermark", os.path.basename(file)))
+
+    print(f"Created zip file of processed images: {output_zip_path}")
     # Return the images and the zip file path
     return processed_images, output_zip_path
 
-def gradio_interface(zip_file, crop_mode, remove_bg, watermark, output_format, bg_choice):
+def gradio_interface(zip_file, crop_mode, remove_bg, bg_method, watermark, output_format, bg_choice):
+    print("Initializing Gradio interface processing...")
     progress = gr.Progress()  # Initialize progress
     watermark_path = watermark.name if watermark else None
-    return process_images(zip_file.name, crop_mode, remove_bg, watermark_path, output_format, bg_choice, progress)
+    return process_images(zip_file.name, crop_mode, remove_bg, bg_method, watermark_path, output_format, bg_choice, progress)
 
 def show_bg_choice(remove_bg, output_format):
     if remove_bg == 'yes' and output_format == 'PNG':
         return gr.update(visible=True)
     return gr.update(visible=False)
 
+def show_bg_method(remove_bg):
+    if remove_bg == 'yes':
+        return gr.update(visible=True)
+    return gr.update(visible=False)
+
 # Create the Gradio interface
 with gr.Blocks() as iface:
     gr.Markdown("# Image Background Removal and Resizing with Optional Watermark")
@@ -164,21 +191,26 @@ with gr.Blocks() as iface:
     with gr.Row():
         crop_mode = gr.Radio(choices=["center", "top", "bottom", "left", "right"], label="Crop Mode", value="center")
         remove_bg = gr.Radio(choices=["yes", "no"], label="Remove Background", value="yes")
-        output_format = gr.Radio(choices=["PNG", "JPG"], label="Output Format", value="PNG")
+        
+        output_format = gr.Radio(choices=["PNG", "JPG"], label="Output Format", value="PNG")    
+
+    with gr.Row():
+        bg_method = gr.Radio(choices=["bria", "rembg"], label="Background Removal Method", value="bria", visible=True)
         bg_choice = gr.Radio(choices=["transparent", "white"], label="Background Choice", value="transparent", visible=True)
 
     remove_bg.change(show_bg_choice, inputs=[remove_bg, output_format], outputs=bg_choice)
+    remove_bg.change(show_bg_method, inputs=remove_bg, outputs=bg_method)
     output_format.change(show_bg_choice, inputs=[remove_bg, output_format], outputs=bg_choice)
 
     gallery = gr.Gallery(label="Processed Images")
     output_zip = gr.File(label="Download Processed Images as ZIP")
 
-    def process(zip_file, crop_mode, remove_bg, watermark, output_format, bg_choice):
-        processed_images, zip_path = gradio_interface(zip_file, crop_mode, remove_bg, watermark, output_format, bg_choice)
+    def process(zip_file, crop_mode, remove_bg, bg_method, watermark, output_format, bg_choice):
+        processed_images, zip_path = gradio_interface(zip_file, crop_mode, remove_bg, bg_method, watermark, output_format, bg_choice)
         return processed_images, zip_path
 
     process_button = gr.Button("Process Images")
-    process_button.click(process, inputs=[zip_file, crop_mode, remove_bg, watermark, output_format, bg_choice], outputs=[gallery, output_zip])
+    process_button.click(process, inputs=[zip_file, crop_mode, remove_bg, bg_method, watermark, output_format, bg_choice], outputs=[gallery, output_zip])
 
 # Launch the interface
 iface.launch()

requirements.txt

@@ -1,3 +1,10 @@
 gradio==4.37.2
-Pillow==10.4.0
 rembg==2.0.57
+torch==2.3.1
+torchvision==0.18.1
+pillow==10.4.0
+numpy==1.26.4
+typing==3.7.4.3
+scikit-image==0.24.0
+huggingface_hub==0.23.4
+transformers==4.42.4