app.py

import os, re
import zipfile
import shutil
import time
from PIL import Image, ImageDraw, ImageFont
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

class LoadModel:
  def __init__(self):
    self.device = "cuda" if torch.cuda.is_available() else "cpu"

  def get_stable_diffusion_model(self):
    return StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", torch_dtype=torch.float32).to(self.device)

  def get_bria_model(self):
    return pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True, device=self.device)

  def remove_background_bria(self, input_path, pipeline):
    print(f"Removing background using bria for image: {input_path}")
    result = pipeline(input_path)
    return result

  def remove_background_rembg(self, 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

class ImageProcessor:
    def __init__(self):
        self.model_loader = LoadModel()
        self.bria_pipeline = self.model_loader.get_bria_model()

    def get_bounding_box_with_threshold(self, 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
        top, bottom = np.where(rows)[0][[0, -1]]
        left, right = np.where(cols)[0][[0, -1]]

        if left < right and top < bottom:
            return (left, top, right, bottom)
        else:
            return None

    def position_logic(self, image_path, canvas_size, padding_top, padding_right, padding_bottom, padding_left, use_threshold=True):
        print('masuk ke method position_logic')
        image = Image.open(image_path)
        image = image.convert("RGBA")
        print('image converted to RGBA')

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

        if bbox:
            # Check 1 pixel around the image for non-transparent pixels
            width, height = image.size
            cropped_sides = []

            # Define tolerance for transparency
            tolerance = 30  # Adjust this value as needed

            # Check top edge
            if any(image.getpixel((x, 0))[3] > tolerance for x in range(width)):
                cropped_sides.append("top")

            # Check bottom edge
            if any(image.getpixel((x, height - 1))[3] > tolerance for x in range(width)):
                cropped_sides.append("bottom")

            # Check left edge
            if any(image.getpixel((0, y))[3] > tolerance for y in range(height)):
                cropped_sides.append("left")

            # Check right edge
            if any(image.getpixel((width - 1, y))[3] > tolerance for y in range(height)):
                cropped_sides.append("right")

            if cropped_sides:
                info_message = f"Info for {os.path.basename(image_path)}: The following sides of the image may contain cropped objects: {', '.join(cropped_sides)}"
                print(info_message)
                log.append({"info": info_message})
            else:
                info_message = f"Info for {os.path.basename(image_path)}: The image is not cropped."
                print(info_message)
                print("ini nih cropped side",cropped_sides)
                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 the new size to expand the image
            target_width, target_height = canvas_size
            aspect_ratio = image.width / image.height

            if len(cropped_sides) == 4:
                # If the image is cropped on all sides, center crop it to fit the canvas
                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)
                    image = image.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)
                    image = image.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
                print(cropped_sides)
            elif not cropped_sides:
                # If the image is not cropped, expand it from center until it touches the padding
                new_height = target_height - padding_top - padding_bottom
                new_width = int(new_height * aspect_ratio)

                if new_width > target_width - padding_left - padding_right:
                    # If width exceeds available space, adjust based on width
                    new_width = target_width - padding_left - padding_right
                    new_height = int(new_width / aspect_ratio)

                # Resize the image
                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:
                # Logic for handling cropped images
                new_height = target_height - padding_bottom
                new_width = int(new_height * aspect_ratio)

                # If new width exceeds canvas width, adjust based on width
                if new_width > target_width:
                    new_width = target_width
                    new_height = int(new_width / aspect_ratio)

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

                # Set position
                if "left" in cropped_sides:
                    x = 0
                else:
                    x = target_width - new_width
                y = 0

        return log, image, x, y

    def process_single_image(self, image_path, output_folder, bg_method, canvas_size_name, output_format, bg_choice, custom_color, watermark_path=None):
        print('masuk ke method process_single_image')
        add_padding_line = False

        if canvas_size_name == 'Rox':
            canvas_size = (1080, 1080)
            padding_top = 112
            padding_right = 125
            padding_bottom = 116
            padding_left = 125
        elif canvas_size_name == 'Columbia':
            canvas_size = (730, 610)
            padding_top = 30
            padding_right = 105
            padding_bottom = 35
            padding_left = 105
        elif canvas_size_name == 'Zalora':
            canvas_size = (763, 1100)
            padding_top = 50
            padding_right = 50
            padding_bottom = 200
            padding_left = 50


        filename = os.path.basename(image_path)
        try:
            print(f"Processing image: {filename}")
            if bg_method == 'rembg':
                image_with_no_bg = self.model_loader.remove_background_rembg(image_path)
            elif bg_method == 'bria':
                image_with_no_bg = self.model_loader.remove_background_bria(image_path, self.bria_pipeline)
            elif bg_method == None:
                image_with_no_bg = Image.open(image_path)

            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 = self.position_logic(temp_image_path, canvas_size, padding_top, padding_right, padding_bottom, padding_left)

            # Create a new canvas with the appropriate background
            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 visible black line for padding when background is not transparent
            if add_padding_line:
                draw = ImageDraw.Draw(canvas)
                draw.rectangle([padding_left, padding_top, canvas_size[0] - padding_right, canvas_size[1] - padding_bottom], outline="black", width=5)
                log.append({"action": "add_padding_line"})

            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)

            # Apply watermark only if the filename ends with "_01" and watermark_path is provided
            if os.path.splitext(filename)[0].endswith("_01") and watermark_path:
                watermark = Image.open(watermark_path).convert("RGBA")
                canvas = canvas.convert("RGBA")
                canvas.paste(watermark, (0, 0), watermark)
                log.append({"action": "add_watermark"})

            if output_format == 'JPG':
                canvas = canvas.convert('RGB')
                canvas.save(output_path, format='JPEG')
            else:
                canvas.save(output_path, format='PNG')

            os.remove(temp_image_path)

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

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

    def remove_extension(self, filename):
        # Regular expression to match any extension at the end of the string
        return re.sub(r'\.[^.]+$', '', filename)

    def process_images(self, 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()):
        print('masuk ke method process_images')
        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 = []

        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)
            except zipfile.BadZipFile as e:
                print(f"Error extracting zip file: {e}")
                return [], None, 0

            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'))]
        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)
        print(f"Total images to process: {total_images}")

        avg_processing_time = 0
        with ThreadPoolExecutor(max_workers=num_workers) as executor:
            future_to_image = {executor.submit(self.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:
                        if watermark_path:
                            get_name = future_to_image[future].split('/')
                            get_name = self.remove_extension(get_name[len(get_name)-1])
                            twibbon_input = f'{get_name}.png' if output_format == 'PNG' else f'{get_name}.jpg'
                            twibbon_output_path = os.path.join(output_folder, f'result_{start_time_image}.png')
                            print('mencoba method add twibon')
                            self.add_twibbon(f'processed_images/{twibbon_input}', watermark_path, twibbon_output_path)
                            print('method add_twibon berhasil')
                            processed_images.append((twibbon_output_path, twibbon_output_path))
                        else:
                            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:
                    print('HAYOLOH KETAUAN ERRORNYA')
                    print(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)
        print("Comprehensive log saved to", os.path.join(output_folder, 'process_log.json'))

        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 remove_white_background(self, twibbon, tolerance=100):
        """
        Menghapus background putih dengan toleransi tertentu.
        tolerance: Nilai antara 0 (tidak toleran, hanya putih murni) hingga 255 (sangat toleran, mencakup hampir semua warna cerah).
        """
        twibbon = twibbon.convert("RGBA")
        data = twibbon.getdata()

        new_data = []
        for item in data:
            # Hitung jarak warna ke putih (255, 255, 255)
            distance_to_white = sum([abs(255 - c) for c in item[:3]])  # RGB distance

            if distance_to_white <= tolerance:
                # Jika jarak warna ke putih lebih kecil dari toleransi, buat transparan
                new_data.append((255, 255, 255, 0))  # Transparan penuh
            else:
                # Tetap pertahankan warna asli
                new_data.append(item)

        twibbon.putdata(new_data)
        return twibbon

    def adjust_opacity(self, twibbon, opacity_level):
        twibbon = twibbon.convert("RGBA")
        data = twibbon.getdata()

        new_data = []
        for item in data:
            # Ubah hanya nilai alpha (transparansi)
            new_alpha = int(item[3] * opacity_level / 255)  # Sesuaikan alpha sesuai opacity_level
            new_data.append((item[0], item[1], item[2], new_alpha))

        twibbon.putdata(new_data)
        return twibbon


    def add_twibbon(self, image_path, twibbon_path, output_path):
        # Open the original image
        image = Image.open(image_path).convert("RGBA")
        print('Original image loaded')

        # Open the twibbon (watermark)
        twibbon = Image.open(twibbon_path).convert("RGBA")
        print('Twibbon (watermark) loaded')

        # Remove white background from twibbon
        twibbon = self.remove_white_background(twibbon)
        # twibbon = self.adjust_opacity(twibbon, 128)

        # Resize the twibbon (watermark)
        image_width, image_height = image.size
        twibbon_size = (image_width // 5, image_height // 5)  # Resize twibbon to 20% of image size
        twibbon = twibbon.resize(twibbon_size, Image.Resampling.LANCZOS)

        # Center the watermark
        twibbon_width, twibbon_height = twibbon.size
        x_offset = (image_width - twibbon_width) // 2
        y_offset = (image_height - twibbon_height) // 2

        # Create a new transparent layer for the watermark
        transparent_layer = Image.new("RGBA", (image_width, image_height), (0, 0, 0, 0))
        transparent_layer.paste(twibbon, (x_offset, y_offset), mask=twibbon.split()[3])

        # Composite the image with the transparent layer
        final_image = Image.alpha_composite(image, transparent_layer)

        # Save the final result
        print('Saving the final image with watermark...')
        final_image.save(output_path)
        print('Image saved successfully')

        return final_image

class ModelInference:
  def __init__(self):
      self.loader = LoadModel()
      self.sd_model = self.loader.get_stable_diffusion_model()

  def text_to_image(self, prompt):
      os.makedirs("generated_images", exist_ok=True)  # Ensure the directory exists
      image = self.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(self, 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 = self.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

class CreativeImageSuite:
    def __init__(self):
        self.inference = ModelInference()  # Use the ModelInference class
        self.processor = ImageProcessor()
        self.theme = "NoCrypt/[email protected]"
        self.title = "# 🎨 Creative Image Suite: Generate, Modify, and Enhance Your Visuals"
        self.description = """
        **Unlock your creativity with our comprehensive image processing tool! This suite offers three powerful features:**
        1. **✏️ Text to Image**: Transform your ideas into stunning visuals by simply entering a descriptive text prompt.
        2. **🖼️ Image to Image**: Enhance existing images by providing a text description of the modifications you want.
        3. **🖌️ Image Background Removal and Resizing**: Effortlessly remove backgrounds from images, resize them, and even add watermarks (optional).
        """
        self.interface = None

    def gradio_interface(self, input_files, bg_method, watermark, canvas_size, output_format, bg_choice, custom_color, num_workers):
        progress = gr.Progress()
        watermark_path = watermark.name if watermark else None

        # Check input_files, is it single image, list image, or zip/rar
        if isinstance(input_files, str) and input_files.lower().endswith(('.zip', '.rar')):
                return self.processor.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 self.processor.process_images(input_files, bg_method, watermark_path, canvas_size, output_format, bg_choice, custom_color, num_workers, progress)
        else:
            return self.processor.process_images(input_files.name, bg_method, watermark_path, canvas_size, output_format, bg_choice, custom_color, num_workers, progress)

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

        return gr.update(visible=False)

    def update_compare(self, evt: gr.SelectData):
        if isinstance(evt.value, dict) and 'caption' in evt.value:
            input_path = evt.value['caption']
            output_path = evt.value['image']['path']
            input_path = input_path.split("Input: ")[-1]
            # 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:
            print("No caption found in selection")
            return gr.update(value=None), gr.update(value=None), gr.update(value=None), gr.update(value=None)

    def master_process(self, input_files, bg_method, watermark, canvas_size, output_format, bg_choice, custom_color, num_workers):
        _, processed_images, zip_path, time_taken = self.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"

    def build_interface(self):
        with gr.Blocks(theme=self.theme) as self.interface:
            # App Header
            gr.Markdown(self.title)
            gr.Markdown(self.description)

            # Text-to-Image Section
            gr.Markdown("## Text to Image Feature")
            gr.Markdown("""
            *Create Visuals from Your Imagination*

            This feature allows you to generate unique images from a simple text description. Just type your idea in the box, and our tool will bring it to life in seconds.  
            - Example Prompts:
                - "Generate an image of a kitchenset with modern furniture."
                - "Generate an image of a casual turtleneck shirt."
                - "A futuristic laptop on a white background ."
            """)
            with gr.Row():
                prompt_input = gr.Textbox(label="Describe your vision to generate a unique image:")
                generate_button = gr.Button("Create Image")
                output_image = gr.Image(label="Your Generated Image")
                download_button = gr.File(label="Download Image", type="filepath")

                generate_button.click(self.inference.text_to_image,
                                      inputs=prompt_input,
                                      outputs=[output_image, download_button])

            # Image-to-Image Section
            gr.Markdown("## Image to Image Feature")
            gr.Markdown("""
            *Enhance or Transform Your Existing Images*

            Upload an image and describe the changes you'd like to see. From subtle edits to artistic transformations, this feature helps you bring new life to your visuals.  
            - Example Edits:
                - "Change the Color of the bag"
                - "Add plastic to the suitcase."
                - "mirror this photo of denim jacket."
            """)
            with gr.Row():
                input_image = gr.Image(label="Upload an image to modify:", type="pil")
                prompt_modification = gr.Textbox(label="Describe the changes you want:")
                modify_button = gr.Button("Apply Changes")
                modified_output_image = gr.Image(label="Your Modified Image")
                download_modified_button = gr.File(label="Download Modified Image", type="filepath")

                modify_button.click(self.inference.text_image_to_image,
                                    inputs=[input_image, prompt_modification],
                                    outputs=[modified_output_image, download_modified_button])

            # Background Removal and Resizing Section
            gr.Markdown("## Image Background Removal and Resizing with Optional Watermark")
            gr.Markdown("""
            *Perfect Your Images for Any Use Case*

            Easily remove backgrounds, resize images to fit your needs, and even add watermarks to maintain originality or branding. This feature is ideal for e-commerce, social media, and design projects.  
            - Features:
                - Supports batch processing of multiple images or ZIP/RAR files.
                - Options for transparent, solid color, or custom backgrounds.
                - Output in your choice of PNG or JPG format.
            """)
            with gr.Row():
                input_files = gr.File(label="Upload an image or a ZIP/RAR file for batch processing:", file_types=[".zip", ".rar", "image"], interactive=True)
                watermark = gr.File(label="Upload an optional watermark (PNG only):", file_types=[".png"])

            with gr.Row():
                canvas_size = gr.Radio(choices=["Rox", "Columbia", "Zalora"], label="Select the desired canvas size:", value="Rox")
                output_format = gr.Radio(choices=["PNG", "JPG"], label="Choose the output format:", value="JPG")
                num_workers = gr.Slider(minimum=1, maximum=16, step=1, label="Set the number of processing threads:", value=5)

            with gr.Row():
                bg_method = gr.Radio(choices=["bria", "rembg", None], label="Choose a background removal method:", value="bria")
                bg_choice = gr.Radio(choices=["transparent", "white", "custom"], label="Select a background style:", value="white")
                custom_color = gr.ColorPicker(label="Pick a custom background color (if applicable):", value="#ffffff", visible=False)

            process_button = gr.Button("Start Processing")
            with gr.Row():
                gallery_processed = gr.Gallery(label="Processed Images")
            with gr.Row():
                image_original = gr.Image(label="Original Image Preview", interactive=False)
                image_processed = gr.Image(label="Processed Image Preview", interactive=False)
            with gr.Row():
                original_ratio = gr.Textbox(label="Aspect Ratio (Original)")
                processed_ratio = gr.Textbox(label="Aspect Ratio (Processed)")
            with gr.Row():
                output_zip = gr.File(label="Download All Processed Images (ZIP)")
                processing_time = gr.Textbox(label="Total Processing Time (seconds)")

            bg_choice.change(self.show_color_picker, inputs=bg_choice, outputs=custom_color)
            process_button.click(self.master_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(self.update_compare, outputs=[image_original, image_processed, original_ratio, processed_ratio])

    def launch(self):
        if self.interface is None:
            self.build_interface()
        self.interface.launch(share=True, debug=True)

app = CreativeImageSuite()
app.launch()