app.py

import os
import zipfile
import shutil
import time
from PIL import Image, ImageDraw
import io
from rembg import remove
import gradio as gr
from concurrent.futures import ThreadPoolExecutor
from diffusers import StableDiffusionPipeline
from transformers import pipeline
import numpy as np
import json
import torch
import logging

# Load Stable Diffusion Model
def load_stable_diffusion_model():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", torch_dtype=torch.float32).to(device)
    return pipe

# Initialize the model globally
sd_model = load_stable_diffusion_model()

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):
    """Remove background using the Bria model."""
    print("Removing background using Bria for the image.")
    
    # Create a segmentation pipeline
    pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True)

    # Load the image
    input_image = Image.open(input_path).convert("RGBA")
    
    # Get the segmentation output
    pillow_mask = pipe(input_image, return_mask=True)  # Outputs a pillow mask
    print("Mask obtained:", pillow_mask)  # Debugging output

    # Create an output image based on the mask
    output_image = Image.new("RGBA", input_image.size)

    # Use the mask to create the output image
    for x in range(input_image.width):
        for y in range(input_image.height):
            # Assuming mask is in a binary format where foreground is True
            if pillow_mask.getpixel((x, y)) > 0:  # Adjust based on actual mask values
                output_image.putpixel((x, y), input_image.getpixel((x, y)))
            else:
                output_image.putpixel((x, y), (0, 0, 0, 0))  # Set to transparent

    return output_image

# Fungsi untuk memproses gambar menggunakan prompt
def text_to_image(prompt):
    os.makedirs("generated_images", exist_ok=True)  # Ensure the directory exists
    image = sd_model(prompt).images[0]  # Generate image using the model
    # Create a sanitized filename by replacing spaces with underscores
    image_path = f"generated_images/{prompt.replace(' ', '_')}.png"
    image.save(image_path)  # Save the generated image
    return image, image_path  # Return the image and its path

# Function to modify an image based on a text prompt
def text_image_to_image(input_image, prompt):
    os.makedirs("generated_images", exist_ok=True)  # Ensure the directory exists
    # Convert input image to PIL Image if necessary
    if not isinstance(input_image, Image.Image):
        input_image = Image.open(input_image)  # Load image from path if given as string
    # Generate modified image using the model with the input image and prompt
    modified_image = sd_model(prompt, init_image=input_image, strength=0.75).images[0]
    # Create a sanitized filename for the modified image
    image_path = f"generated_images/{prompt.replace(' ', '_')}_modified.png"
    modified_image.save(image_path)  # Save the modified image
    return modified_image, image_path  # Return the modified image and its path

def get_bounding_box_with_threshold(image, threshold):
    # Convert image to numpy array
    img_array = np.array(image)
    
    # 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
    if np.any(rows) and np.any(cols):
        top, bottom = np.where(rows)[0][[0, -1]]
        left, right = np.where(cols)[0][[0, -1]]
        return (left, top, right, bottom)
    else:
        return None
        
def position_logic(image_path, canvas_size, padding_top, padding_right, padding_bottom, padding_left, use_threshold=True):
    """Position and resize an image based on cropping and padding requirements."""
    image = Image.open(image_path).convert("RGBA")

    # Get the bounding box of the non-blank area with threshold
    bbox = get_bounding_box_with_threshold(image, threshold=10) if use_threshold else image.getbbox()
    log = []

    if bbox:
        width, height = image.size
        cropped_sides = []
        tolerance = 30  # Define a constant for transparency tolerance

        # Check each edge for non-transparent pixels
        for edge in ['top', 'bottom', 'left', 'right']:
            if edge == 'top' and any(image.getpixel((x, 0))[3] > tolerance for x in range(width)):
                cropped_sides.append(edge)
            elif edge == 'bottom' and any(image.getpixel((x, height - 1))[3] > tolerance for x in range(width)):
                cropped_sides.append(edge)
            elif edge == 'left' and any(image.getpixel((0, y))[3] > tolerance for y in range(height)):
                cropped_sides.append(edge)
            elif edge == 'right' and any(image.getpixel((width - 1, y))[3] > tolerance for y in range(height)):
                cropped_sides.append(edge)

        # Log cropping information
        info_message = f"Info for {os.path.basename(image_path)}: The following sides of the image may contain cropped objects: {', '.join(cropped_sides) if cropped_sides else 'none'}"
        print(info_message)
        log.append({"info": info_message})

        # Crop the image to the bounding box
        image = image.crop(bbox)
        log.append({"action": "crop", "bbox": [str(bbox[0]), str(bbox[1]), str(bbox[2]), str(bbox[3])]})

        # Calculate new size to expand the image
        target_width, target_height = canvas_size
        aspect_ratio = image.width / image.height

        # Handling image positioning and resizing based on cropped sides
        if len(cropped_sides) == 4:
            # Center crop if cropped on all sides
            if aspect_ratio > 1:  # Landscape
                new_height = target_height
                new_width = int(new_height * aspect_ratio)
                left = (new_width - target_width) // 2
                image = image.resize((new_width, new_height), Image.LANCZOS).crop((left, 0, left + target_width, target_height))
            else:  # Portrait or square
                new_width = target_width
                new_height = int(new_width / aspect_ratio)
                top = (new_height - target_height) // 2
                image = image.resize((new_width, new_height), Image.LANCZOS).crop((0, top, target_width, top + target_height))
            log.append({"action": "center_crop_resize", "new_size": f"{target_width}x{target_height}"})
            x, y = 0, 0
        elif not cropped_sides:
            # Expand from center
            new_height = target_height - padding_top - padding_bottom
            new_width = int(new_height * aspect_ratio)
            if new_width > target_width - padding_left - padding_right:
                new_width = target_width - padding_left - padding_right
                new_height = int(new_width / aspect_ratio)

            image = image.resize((new_width, new_height), Image.LANCZOS)
            log.append({"action": "resize", "new_width": str(new_width), "new_height": str(new_height)})
            x = (target_width - new_width) // 2
            y = target_height - new_height - padding_bottom
        else:
            # Handle specific cropped side cases
            new_width, new_height = target_width, target_height  # Default values for resizing
            for side in cropped_sides:
                if side in ["top", "bottom"]:
                    new_height = target_height - (padding_top + padding_bottom)
                if side in ["left", "right"]:
                    new_width = target_width - (padding_left + padding_right)

            image = image.resize((new_width, new_height), Image.LANCZOS)
            log.append({"action": "resize", "new_width": str(new_width), "new_height": str(new_height)})

            # Determine position based on cropped sides
            x = (target_width - new_width) // 2 if "left" not in cropped_sides and "right" not in cropped_sides else (0 if "left" in cropped_sides else target_width - new_width)
            y = (0 if "top" in cropped_sides else target_height - new_height)

            log.append({"action": "position", "x": str(x), "y": str(y)})

    return log, image, x, y

# Constants for canvas sizes and paddings
CANVAS_SIZES = {
    'Rox': ((1080, 1080), (112, 125, 116, 125)),
    'Columbia': ((730, 610), (30, 105, 35, 105)),
    'Zalora': ((763, 1100), (50, 50, 200, 50))
}

def process_single_image(image_path, output_folder, bg_method, canvas_size_name, output_format, bg_choice, custom_color, watermark_path=None):
    """Processes a single image by removing its background and applying various transformations.

    Args:
        image_path (str): Path to the input image.
        output_folder (str): Path to the output folder.
        bg_method (str): Background removal method ('rembg' or 'bria').
        canvas_size_name (str): Name of the canvas size.
        output_format (str): Desired output format ('JPG' or 'PNG').
        bg_choice (str): Background choice ('white', 'custom', or 'transparent').
        custom_color (tuple): Custom background color as an RGBA tuple.
        watermark_path (str, optional): Path to a watermark image.

    Returns:
        tuple: A tuple containing the output path and log of actions performed.
    """
    try:
        canvas_size, (padding_top, padding_right, padding_bottom, padding_left) = CANVAS_SIZES[canvas_size_name]
        filename = os.path.basename(image_path)
        logging.info(f"Processing image: {filename}")

        # Remove background
        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)
        else:
            raise ValueError("Invalid background method specified.")

        # Temporary file for processed image
        temp_image_path = os.path.join(output_folder, f"temp_{filename}")
        image_with_no_bg.save(temp_image_path, format='PNG')

        log, new_image, x, y = position_logic(temp_image_path, canvas_size, padding_top, padding_right, padding_bottom, padding_left)

        # Create canvas based on background choice
        if bg_choice == 'white':
            canvas = Image.new("RGBA", canvas_size, "WHITE")
        elif bg_choice == 'custom':
            canvas = Image.new("RGBA", canvas_size, custom_color)
        else:  # transparent
            canvas = Image.new("RGBA", canvas_size, (0, 0, 0, 0))

        # Paste the resized image onto the canvas
        canvas.paste(new_image, (x, y), new_image)
        log.append({"action": "paste", "position": [str(x), str(y)]})

        # Add watermark if applicable
        if watermark_path:
            watermark = Image.open(watermark_path).convert("RGBA")
            # Hitung posisi untuk menempatkan watermark di tengah
            watermark_width, watermark_height = watermark.size
            canvas_width, canvas_height = canvas.size
            # Hitung posisi (x, y) untuk watermark
            x = (canvas_width - watermark_width) // 2
            y = (canvas_height - watermark_height) // 2
            # Tempelkan watermark
            canvas.paste(watermark, (x, y), watermark)
            log.append({"action": "add_watermark"})

        # Determine output format
        output_ext = 'jpg' if output_format == 'JPG' else 'png'
        output_filename = f"{os.path.splitext(filename)[0]}.{output_ext}"
        output_path = os.path.join(output_folder, output_filename)

        # Save the canvas in the desired format
        if output_format == 'JPG':
            canvas.convert('RGB').save(output_path, format='JPEG')
        else:
            canvas.save(output_path, format='PNG')

        # Clean up temporary file
        os.remove(temp_image_path)

        logging.info(f"Processed image path: {output_path}")
        return [(output_path, image_path)], log

    except Exception as e:
        logging.error(f"Error processing {filename}: {e}")
        return None, None
    
# Set up logging
logging.basicConfig(level=logging.INFO)

def process_images(input_files, bg_method='rembg', watermark_path=None, canvas_size='Rox', output_format='PNG', bg_choice='transparent', custom_color="#ffffff", num_workers=4, progress=gr.Progress()):
    """Processes images by removing backgrounds and applying various transformations.

    Args:
        input_files (str or list): Path to a ZIP file or a list of image paths.
        bg_method (str): Background removal method ('rembg' or 'bria').
        watermark_path (str, optional): Path to a watermark image.
        canvas_size (str): Name of the canvas size.
        output_format (str): Desired output format ('PNG' or 'JPG').
        bg_choice (str): Background choice ('transparent', 'white', or 'custom').
        custom_color (str): Custom background color in hex format.
        num_workers (int): Number of parallel workers for processing.
        progress (gr.Progress): Progress tracking interface.

    Returns:
        tuple: A tuple containing original images, processed images, output zip path, and total processing time.
    """
    start_time = time.time()
    
    output_folder = "processed_images"
    if os.path.exists(output_folder):
        shutil.rmtree(output_folder)
    os.makedirs(output_folder)

    processed_images = []
    original_images = []
    all_logs = []
    
    image_files = []

    if isinstance(input_files, str) and input_files.lower().endswith(('.zip', '.rar')):
        # Handle zip file
        input_folder = "temp_input"
        if os.path.exists(input_folder):
            shutil.rmtree(input_folder)
        os.makedirs(input_folder)

        try:
            with zipfile.ZipFile(input_files, 'r') as zip_ref:
                zip_ref.extractall(input_folder)
            image_files = [os.path.join(input_folder, f) for f in os.listdir(input_folder) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif', '.webp'))]
        except zipfile.BadZipFile as e:
            logging.error(f"Error extracting zip file: {e}")
            return [], None, 0
    elif isinstance(input_files, list):
        # Handle multiple files
        image_files = input_files
    else:
        # Handle single file
        image_files = [input_files]
    
    total_images = len(image_files)
    logging.info(f"Total images to process: {total_images}")

    avg_processing_time = 0
    with ThreadPoolExecutor(max_workers=num_workers) as executor:
        future_to_image = {executor.submit(process_single_image, image_path, output_folder, bg_method, canvas_size, 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:
                start_time_image = time.time()
                result, log = future.result()
                end_time_image = time.time()
                image_processing_time = end_time_image - start_time_image
                
                # Update average processing time
                avg_processing_time = (avg_processing_time * idx + image_processing_time) / (idx + 1)
                
                if result:
                    processed_images.extend(result)
                    original_images.append(future_to_image[future])
                    all_logs.append({os.path.basename(future_to_image[future]): log})
                
                # Estimate remaining time
                remaining_images = total_images - (idx + 1)
                estimated_remaining_time = remaining_images * avg_processing_time
                
                progress((idx + 1) / total_images, f"{idx + 1}/{total_images} images processed. Estimated time remaining: {estimated_remaining_time:.2f} seconds")
            except Exception as e:
                logging.error(f"Error processing image {future_to_image[future]}: {e}")

    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))

    # Write the comprehensive log for all images
    with open(os.path.join(output_folder, 'process_log.json'), 'w') as log_file:
        json.dump(all_logs, log_file, indent=4)
    logging.info("Comprehensive log saved to %s", os.path.join(output_folder, 'process_log.json'))

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

# Set up logging
logging.basicConfig(level=logging.INFO)

def gradio_interface(input_files, bg_method, watermark, canvas_size, output_format, bg_choice, custom_color, num_workers):
    """Handles input files and processes them accordingly."""
    progress = gr.Progress()
    watermark_path = watermark.name if watermark else None
    
    # Check the type of input_files and call the process_images function
    if isinstance(input_files, str) and input_files.lower().endswith(('.zip', '.rar')):
        return process_images(input_files, bg_method, watermark_path, canvas_size, output_format, bg_choice, custom_color, num_workers, progress)
    elif isinstance(input_files, list):
        return process_images(input_files, bg_method, watermark_path, canvas_size, output_format, bg_choice, custom_color, num_workers, progress)
    else:
        return process_images(input_files.name, bg_method, watermark_path, canvas_size, output_format, bg_choice, custom_color, num_workers, progress)

def show_color_picker(bg_choice):
    """Shows the color picker if 'custom' background is selected."""
    return gr.update(visible=(bg_choice == 'custom'))

def update_compare(evt: gr.SelectData):
    """Updates the displayed images and their ratios when a processed image is selected."""
    if isinstance(evt.value, dict) and 'caption' in evt.value:
        input_path = evt.value['caption'].split("Input: ")[-1]
        output_path = evt.value['image']['path']
        
        # Open the original and processed images
        original_img = Image.open(input_path)
        processed_img = Image.open(output_path)
        
        # Calculate the aspect ratios
        original_ratio = f"{original_img.width}x{original_img.height}"
        processed_ratio = f"{processed_img.width}x{processed_img.height}"
        
        return (gr.update(value=input_path), 
                gr.update(value=output_path), 
                gr.update(value=original_ratio), 
                gr.update(value=processed_ratio)
               )
    else:
        logging.warning("No caption found in selection")
        return (gr.update(value=None),) * 4

def process(input_files, bg_method, watermark, canvas_size, output_format, bg_choice, custom_color, num_workers):
    """Processes the images and returns the results."""
    try:
        _, processed_images, zip_path, time_taken = gradio_interface(input_files, bg_method, watermark, canvas_size, output_format, bg_choice, custom_color, num_workers)
        processed_images_with_captions = [(img, f"Input: {caption}") for img, caption in processed_images]
        return processed_images_with_captions, zip_path, f"{time_taken:.2f} seconds"
    except Exception as e:
        logging.error(f"Error in processing images: {e}")
        return [], None, "Error during processing"

with gr.Blocks(theme="NoCrypt/[email protected]") as iface:
    gr.Markdown("# Image Background Removal and Resizing with Optional Watermark")
    gr.Markdown("Choose to upload multiple images or a ZIP/RAR file, select the crop mode, optionally upload a watermark image, and choose the output format.")
    
    # Fitur Text to Image
    gr.Markdown("# Text to Image Feature")
    with gr.Row():
        prompt_input = gr.Textbox(label="Enter your prompt for image generation:")
        generate_button = gr.Button("Generate Image")
        output_image = gr.Image(label="Generated Image")
        download_button = gr.File(label="Download Generated Image", type="filepath")
        
        generate_button.click(text_to_image, inputs=prompt_input, outputs=[output_image, download_button])

    # Fitur Text Image to Image
    gr.Markdown("# Text Image to Image Feature")
    with gr.Row():
        input_image = gr.Image(label="Upload Image for Modification", type="pil")
        prompt_modification = gr.Textbox(label="Enter your prompt for modification:")
        modify_button = gr.Button("Modify Image")
        modified_output_image = gr.Image(label="Modified Image")
        download_modified_button = gr.File(label="Download Modified Image", type="filepath")
        
        modify_button.click(text_image_to_image, inputs=[input_image, prompt_modification], outputs=[modified_output_image, download_modified_button])
    
    with gr.Row():
        input_files = gr.File(label="Upload Image or ZIP/RAR file", file_types=[".zip", ".rar", "image"], interactive=True)
        watermark = gr.File(label="Upload Watermark Image (Optional)", file_types=[".png"])
    
    with gr.Row():
        canvas_size = gr.Radio(choices=["Rox", "Columbia", "Zalora"], label="Canvas Size", value="Rox")
        output_format = gr.Radio(choices=["PNG", "JPG"], label="Output Format", value="JPG")
        num_workers = gr.Slider(minimum=1, maximum=16, step=1, label="Number of Workers", value=5)

    with gr.Row():
        bg_method = gr.Radio(choices=["bria", "rembg"], label="Background Removal Method", value="bria")
        bg_choice = gr.Radio(choices=["transparent", "white", "custom"], label="Background Choice", value="white")
        custom_color = gr.ColorPicker(label="Custom Background Color", value="#ffffff", visible=False)

    process_button = gr.Button("Process Images")

    with gr.Row():
        gallery_processed = gr.Gallery(label="Processed Images")
    with gr.Row():
        image_original = gr.Image(label="Original Images", interactive=False)
        image_processed = gr.Image(label="Processed Images", interactive=False)
    with gr.Row():
        original_ratio = gr.Textbox(label="Original Ratio")
        processed_ratio = gr.Textbox(label="Processed Ratio")
    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)
    process_button.click(process, inputs=[input_files, bg_method, watermark, canvas_size, output_format, bg_choice, custom_color, num_workers], outputs=[gallery_processed, output_zip, processing_time])
    gallery_processed.select(update_compare, outputs=[image_original, image_processed, original_ratio, processed_ratio])

iface.launch(share=True)