add watermark and output selection

app.py

@@ -49,7 +49,7 @@ def remove_background(input_path):
     img = Image.open(io.BytesIO(output_image)).convert("RGBA")
     return img
 
-def process_single_image(image_path, output_folder, crop_mode, remove_bg):
+def process_single_image(image_path, output_folder, crop_mode, remove_bg, output_format, bg_choice, watermark_path=None):
     filename = os.path.basename(image_path)
     try:
         if remove_bg == 'yes':
@@ -63,9 +63,36 @@ def process_single_image(image_path, output_folder, crop_mode, remove_bg):
         # Resize and crop the image with or without background removal
         new_image = resize_and_crop_image(temp_image_path, crop_mode=crop_mode)
 
+        # Apply watermark if provided
+        if watermark_path:
+            watermark = Image.open(watermark_path).convert("RGBA")
+            new_image.paste(watermark, (0, 0), watermark)
+
         # Save the final image
-        output_path = os.path.join(output_folder, f"{os.path.splitext(filename)[0]}.png")
-        new_image.save(output_path, format='PNG')
+        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}")
+            if output_format == 'JPG':
+                new_image.convert('RGB').save(output_path, format='JPEG')
+            else:
+                new_image.save(output_path, format='PNG')
 
         if remove_bg == 'yes':
             # Remove the temporary file
@@ -77,7 +104,7 @@ def process_single_image(image_path, output_folder, crop_mode, remove_bg):
         print(f"Error processing {filename}: {e}")
         return None
 
-def process_images(zip_file, crop_mode='center', remove_bg='yes', progress=gr.Progress()):
+def process_images(zip_file, crop_mode='center', remove_bg='yes', watermark_path=None, output_format='PNG', bg_choice='transparent', progress=gr.Progress()):
     # Create a temporary directory
     input_folder = "temp_input"
     output_folder = "temp_output"
@@ -98,7 +125,7 @@ def process_images(zip_file, crop_mode='center', remove_bg='yes', progress=gr.Pr
 
     # 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): image_path for image_path in image_files}
+        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}
         for idx, future in enumerate(future_to_image):
             result = future.result()
             if result:
@@ -115,29 +142,43 @@ def process_images(zip_file, crop_mode='center', remove_bg='yes', progress=gr.Pr
     # Return the images and the zip file path
     return processed_images, output_zip_path
 
-def gradio_interface(zip_file, crop_mode, remove_bg):
+def gradio_interface(zip_file, crop_mode, remove_bg, watermark, output_format, bg_choice):
     progress = gr.Progress()  # Initialize progress
-    return process_images(zip_file.name, crop_mode, remove_bg, 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)
+
+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)
 
 # Create the Gradio interface
 with gr.Blocks() as iface:
-    gr.Markdown("# Image Background Removal and Resizing")
-    gr.Markdown("Upload a ZIP or RAR file containing images, choose the crop mode, and get the processed images.")
-    
+    gr.Markdown("# Image Background Removal and Resizing with Optional Watermark")
+    gr.Markdown("Upload a ZIP or RAR file containing images, choose the crop mode, optionally upload a watermark image, and select the output format.")
+
     with gr.Row():
         zip_file = gr.File(label="Upload ZIP/RAR file of images", file_types=[".zip", ".rar"])
+        watermark = gr.File(label="Upload Watermark Image (Optional)", file_types=[".png"])
+    
+    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")
+        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)
+    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):
-        processed_images, zip_path = gradio_interface(zip_file, crop_mode, remove_bg)
+    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)
         return processed_images, zip_path
 
     process_button = gr.Button("Process Images")
-    process_button.click(process, inputs=[zip_file, crop_mode, remove_bg], outputs=[gallery, output_zip])
+    process_button.click(process, inputs=[zip_file, crop_mode, remove_bg, watermark, output_format, bg_choice], outputs=[gallery, output_zip])
 
 # Launch the interface
 iface.launch()