app.py

import os
import zipfile
import shutil
import time
from PIL import Image
import io
from rembg import remove
import gradio as gr
from concurrent.futures import ThreadPoolExecutor
from transformers import pipeline

def colors_within_tolerance(color1, color2, tolerance):
    return all(abs(c1 - c2) <= tolerance for c1, c2 in zip(color1, color2))

def check_border_colors(image_path, tolerance):
    image = Image.open(image_path)
    pixels = image.load()
    
    width, height = image.size

    left_border_color = pixels[0, 0]
    right_border_color = pixels[width - 1, 0]

    for y in range(height):
        if not colors_within_tolerance(pixels[0, y], left_border_color, tolerance):
            return False
        if not colors_within_tolerance(pixels[width - 1, y], right_border_color, tolerance):
            return False

    return True

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
        print(f"Original image size: {width}x{height}")

        scaling_factor = max(target_size[0] / width, target_size[1] / height)

        new_size = (int(width * scaling_factor), int(height * scaling_factor))
        resized_img = img.resize(new_size, Image.LANCZOS)
        print(f"Resized image size: {new_size}")

        if crop_mode == 'center':
            left = (resized_img.width - target_size[0]) / 2
            top = (resized_img.height - target_size[1]) / 2
        elif crop_mode == 'top':
            left = (resized_img.width - target_size[0]) / 2
            top = 0
        elif crop_mode == 'bottom':
            left = (resized_img.width - target_size[0]) / 2
            top = resized_img.height - target_size[1]
        elif crop_mode == 'left':
            left = 0
            top = (resized_img.height - target_size[1]) / 2
        elif crop_mode == 'right':
            left = resized_img.width - target_size[0]
            top = (resized_img.height - target_size[1]) / 2

        right = left + target_size[0]
        bottom = top + target_size[1]

        cropped_img = resized_img.crop((left, top, right, bottom))
        print(f"Cropped image size: {cropped_img.size}")

        return cropped_img

def remove_background_rembg(input_path):
    print(f"Removing background using rembg for image: {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 remove_background_bria(input_path):
    print(f"Removing background using bria for image: {input_path}")
    pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True)
    pillow_image = pipe(input_path)
    return pillow_image

def process_single_image(image_path, output_folder, crop_mode, bg_method, output_format, bg_choice, custom_color, watermark_path=None):
    filename = os.path.basename(image_path)
    try:
        print(f"Processing image: {filename}")
        if bg_method == 'rembg':
            image_with_no_bg = remove_background_rembg(image_path)
        elif bg_method == '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')

        if check_border_colors(temp_image_path, tolerance=50):
            print(f"Border colors are the same for image: {filename}")
            if bg_choice == 'transparent':
                new_image = Image.new("RGBA", (1080, 1080), (255, 255, 255, 0))
            else:
                new_image = Image.new("RGBA", (1080, 1080), custom_color)

            width, height = image_with_no_bg.size
            scaling_factor = min(1080 / width, 1080 / height)
            new_size = (int(width * scaling_factor), int(height * scaling_factor))
            resized_img = image_with_no_bg.resize(new_size, Image.LANCZOS)
            print(f"Resized image size: {new_size}")
            new_image.paste(resized_img, ((1080 - resized_img.width) // 2, (1080 - resized_img.height) // 2))
        else:
            print(f"Border colors are different for image: {filename}")
            new_image = resize_and_crop_image(temp_image_path, crop_mode=crop_mode)

        if bg_choice == 'white':
            new_image = new_image.convert("RGBA")
            white_bg = Image.new("RGBA", new_image.size, "WHITE")
            new_image = Image.alpha_composite(white_bg, new_image)
        elif bg_choice == 'custom':
            new_image = new_image.convert("RGBA")
            custom_bg = Image.new("RGBA", new_image.size, custom_color)
            new_image = Image.alpha_composite(custom_bg, new_image)

        images_paths = []

        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)

        if watermark_path:
            watermark = Image.open(watermark_path).convert("RGBA")
            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_with_watermark.convert('RGB').save(output_path_with_watermark, format='JPEG')
            else:
                new_image_with_watermark.save(output_path_with_watermark, format='PNG')
            images_paths.append(output_path_with_watermark)

        os.remove(temp_image_path)

        print(f"Processed image paths: {images_paths}")
        return images_paths

    except Exception as e:
        print(f"Error processing {filename}: {e}")
        return None

def process_images(zip_file, crop_mode='center', bg_method='rembg', watermark_path=None, output_format='PNG', bg_choice='transparent', custom_color="#ffffff", num_workers=4, progress=gr.Progress()):
    start_time = time.time()
    
    input_folder = "temp_input"
    output_folder = "temp_output"
    if os.path.exists(input_folder):
        shutil.rmtree(input_folder)
    if os.path.exists(output_folder):
        shutil.rmtree(output_folder)
    os.makedirs(input_folder)
    os.makedirs(output_folder)
    
    try:
        with zipfile.ZipFile(zip_file, 'r') as zip_ref:
            zip_ref.extractall(input_folder)
    except zipfile.BadZipFile as e:
        print(f"Error extracting zip file: {e}")
        return [], None, 0
    
    processed_images = []
    original_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"Total images to process: {total_images}")

    with ThreadPoolExecutor(max_workers=num_workers) as executor:
        future_to_image = {executor.submit(process_single_image, image_path, output_folder, crop_mode, bg_method, output_format, bg_choice, custom_color, watermark_path): image_path for image_path in image_files}
        for idx, future in enumerate(future_to_image):
            try:
                result = future.result()
                if result:
                    processed_images.extend(result)
                    original_images.append(future_to_image[future])
            except Exception as e:
                print(f"Error processing image {future_to_image[future]}: {e}")
            progress((idx + 1) / total_images, f"{idx + 1}/{total_images} images processed")

    output_zip_path = "processed_images.zip"
    with zipfile.ZipFile(output_zip_path, 'w') as zipf:
        for file in processed_images:
            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)))

    end_time = time.time()
    processing_time = end_time - start_time
    print(f"Processing time: {processing_time} seconds")
    
    return original_images, processed_images, output_zip_path, processing_time

def gradio_interface(zip_file, crop_mode, bg_method, watermark, output_format, bg_choice, custom_color, num_workers):
    progress = gr.Progress()
    watermark_path = watermark.name if watermark else None
    return process_images(zip_file.name, crop_mode, bg_method, watermark_path, output_format, bg_choice, custom_color, num_workers, progress)

def show_bg_choice(output_format):
    if output_format == 'PNG':
        return gr.update(visible=True)
    return gr.update(visible=False)

def show_color_picker(bg_choice):
    if bg_choice == 'custom':
        return gr.update(visible=True)
    return gr.update(visible=False)

with gr.Blocks() as iface:
    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")
        output_format = gr.Radio(choices=["PNG", "JPG"], label="Output Format", value="PNG")
        num_workers = gr.Slider(minimum=1, maximum=16, step=1, label="Number of Workers", value=2)

    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", "custom"], label="Background Choice", value="transparent", visible=True)
        custom_color = gr.ColorPicker(label="Custom Background Color", value="#ffffff", visible=False)

    with gr.Row():
        gallery_original = gr.Gallery(label="Original Images")
        gallery_processed = gr.Gallery(label="Processed Images")
    with gr.Row():
        output_zip = gr.File(label="Download Processed Images as ZIP")
        processing_time = gr.Textbox(label="Processing Time (seconds)")

    bg_choice.change(show_color_picker, inputs=bg_choice, outputs=custom_color)
    output_format.change(show_bg_choice, inputs=output_format, outputs=bg_choice)

    def process(zip_file, crop_mode, bg_method, watermark, output_format, bg_choice, custom_color, num_workers):
        original_images, processed_images, zip_path, time_taken = gradio_interface(zip_file, crop_mode, bg_method, watermark, output_format, bg_choice, custom_color, num_workers)
        return original_images, processed_images, zip_path, f"{time_taken:.2f} seconds"

    process_button = gr.Button("Process Images")
    process_button.click(process, inputs=[zip_file, crop_mode, bg_method, watermark, output_format, bg_choice, custom_color, num_workers], outputs=[gallery_original, gallery_processed, output_zip, processing_time])

iface.launch()