Update app.py

app.py

@@ -12,7 +12,6 @@ from transformers import pipeline
 import numpy as np
 import json
 import torch
-import logging
 
 # Load Stable Diffusion Model
 def load_stable_diffusion_model():
@@ -32,32 +31,15 @@ def remove_background_rembg(input_path):
     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)
+    print(f"Removing background using bria for image: {input_path}")
+    device = 0 if torch.cuda.is_available() else -1
 
-    # 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
+    # Load the segmentation model
+    pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True, device=device)
 
-    return output_image
+    # Process the image
+    result = pipe(input_path)
+    return result
 
 # Function to process images using prompts
 def text_to_image(prompt):
@@ -82,154 +64,360 @@ def text_image_to_image(input_image, prompt):
     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
+	# Convert image to numpy array
     img_array = np.array(image)
-    
+
     # Get alpha channel
-    alpha = img_array[:, :, 3]
-    
+    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]]
+    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(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")
+    image = Image.open(image_path)
+    image = image.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()
+    if use_threshold:
+        bbox = 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 = []
-        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})
+
+        # 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)
+            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
+        # Calculate the 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 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).crop((left, 0, left + target_width, target_height))
+                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).crop((0, top, target_width, top + target_height))
+                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
         elif not cropped_sides:
-            # Expand from center
+            # 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:
-            # 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)})
+            # New logic for handling cropped top and left, or top and right
+            if set(cropped_sides) == {"top", "left"} or set(cropped_sides) == {"top", "right"}:
+                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:  # right in cropped_sides
+                    x = target_width - new_width
+                y = 0
+            
+                # If the resized image is taller than the canvas minus padding, crop from the bottom
+                if new_height > target_height - padding_bottom:
+                    crop_bottom = new_height - (target_height - padding_bottom)
+                    image = image.crop((0, 0, new_width, new_height - crop_bottom))
+                    new_height = target_height - padding_bottom
+                    log.append({"action": "crop_vertical", "bottom_pixels_removed": str(crop_bottom)})
+            
+                log.append({"action": "position", "x": str(x), "y": str(y)})
+            elif set(cropped_sides) == {"bottom", "left"} or set(cropped_sides) == {"bottom", "right"}:
+                # Handle bottom & left or bottom & right cropped images
+                new_height = target_height - padding_top
+                new_width = int(new_height * aspect_ratio)
+            
+                # If new width exceeds canvas width, adjust based on width
+                if new_width > target_width - padding_left - padding_right:
+                    new_width = target_width - padding_left - padding_right
+                    new_height = int(new_width / aspect_ratio)
+            
+                # Resize the image without cropping or stretching
+                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:  # right in cropped_sides
+                    x = target_width - new_width
+                y = target_height - new_height
+            
+                log.append({"action": "position", "x": str(x), "y": str(y)})
+            elif set(cropped_sides) == {"bottom", "left", "right"}:
+                # Expand the image from the center
+                new_width = target_width
+                new_height = int(new_width / aspect_ratio)
+            
+                if new_height < target_height:
+                    new_height = target_height
+                    new_width = int(new_height * aspect_ratio)
+            
+                image = image.resize((new_width, new_height), Image.LANCZOS)
+            
+                # Crop to fit the canvas
+                left = (new_width - target_width) // 2
+                top = 0
+                image = image.crop((left, top, left + target_width, top + target_height))
+            
+                log.append({"action": "expand_and_crop", "new_size": f"{target_width}x{target_height}"})
+                x, y = 0, 0
+            elif cropped_sides == ["top"]:
+                # New logic for handling only top-cropped images
+                if image.width > image.height:
+                    new_width = target_width
+                    new_height = int(target_width / aspect_ratio)
+                else:
+                    new_height = target_height - padding_bottom
+                    new_width = int(new_height * 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 = 0  # Align to top
+            
+                # Apply padding only to non-cropped sides
+                x = max(padding_left, min(x, target_width - new_width - padding_right))
+            elif cropped_sides in [["right"], ["left"]]:
+                # New logic for handling only right-cropped or left-cropped images
+                if image.width > image.height:
+                    new_width = target_width - max(padding_left, padding_right)
+                    new_height = int(new_width / aspect_ratio)
+                else:
+                    new_height = target_height - padding_top - padding_bottom
+                    new_width = int(new_height * 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)})
+            
+                if cropped_sides == ["right"]:
+                    x = target_width - new_width  # Align to right
+                else:  # cropped_sides == ["left"]
+                    x = 0  # Align to left
+                y = target_height - new_height - padding_bottom  # Respect bottom padding
+            
+                # Ensure top padding is respected
+                if y < padding_top:
+                    y = padding_top
+                
+                log.append({"action": "position", "x": str(x), "y": str(y)})
+            elif set(cropped_sides) == {"left", "right"}:
+                # Logic for handling images cropped on both left and right sides
+                new_width = target_width  # Expand to full width of canvas
+            
+                # Calculate the aspect ratio of the original image
+                aspect_ratio = image.width / image.height
+            
+                # Calculate the new height while maintaining aspect ratio
+                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 horizontal position (always 0 as it spans full width)
+                x = 0
+            
+                # Calculate vertical position to respect bottom padding
+                y = target_height - new_height - padding_bottom
+            
+                # If the resized image is taller than the canvas, crop from the top only
+                if new_height > target_height - padding_bottom:
+                    crop_top = new_height - (target_height - padding_bottom)
+                    image = image.crop((0, crop_top, new_width, new_height))
+                    new_height = target_height - padding_bottom
+                    y = 0
+                    log.append({"action": "crop_vertical", "top_pixels_removed": str(crop_top)})
+                else:
+                    # Align the image to the bottom with padding
+                    y = target_height - new_height - padding_bottom
+            
+                log.append({"action": "position", "x": str(x), "y": str(y)})
+            elif cropped_sides == ["bottom"]:
+                # Logic for handling images cropped on the bottom side
+                # Calculate the aspect ratio of the original image
+                aspect_ratio = image.width / image.height
+            
+                if aspect_ratio < 1:  # Portrait orientation
+                    new_height = target_height - padding_top  # Full height with top padding
+                    new_width = int(new_height * aspect_ratio)
+                
+                    # If the new width exceeds the canvas width, adjust it
+                    if new_width > target_width:
+                        new_width = target_width
+                        new_height = int(new_width / aspect_ratio)
+                else:  # Landscape orientation
+                    new_width = target_width - padding_left - padding_right
+                    new_height = int(new_width / aspect_ratio)
+                
+                    # If the new height exceeds the canvas height, adjust it
+                    if new_height > target_height:
+                        new_height = target_height
+                        new_width = int(new_height * 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 horizontal position (centered)
+                x = (target_width - new_width) // 2
+            
+                # Set vertical position (touching bottom edge for all cases)
+                y = target_height - new_height
+            
+                log.append({"action": "position", "x": str(x), "y": str(y)})
+            else:
+                # Use the original resizing logic for other partially cropped images
+                if image.width > image.height:
+                    new_width = target_width
+                    new_height = int(target_width / aspect_ratio)
+                else:
+                    new_height = target_height
+                    new_width = int(target_height * 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)})
+            
+                # Center horizontally for all images
+                x = (target_width - new_width) // 2
+                y = target_height - new_height - padding_bottom
+            
+                # Adjust positions for cropped sides
+                if "top" in cropped_sides:
+                    y = 0
+                elif "bottom" in cropped_sides:
+                    y = target_height - new_height
+                if "left" in cropped_sides:
+                    x = 0
+                elif "right" in cropped_sides:
+                    x = target_width - new_width
+            
+                # Apply padding only to non-cropped sides, but keep horizontal centering
+                if "left" not in cropped_sides and "right" not in cropped_sides:
+                    x = (target_width - new_width) // 2  # Always center horizontally
+                if "top" not in cropped_sides and "bottom" not in cropped_sides:
+                    y = max(padding_top, min(y, target_height - new_height - padding_bottom))
 
     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.
-    """
+    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:
-        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
-        image_with_no_bg = Image.open(image_path).convert("RGBA")
+        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)
-        elif bg_method == 'none':
-            pass  # Do nothing, keep original image with background
-        else:
-            raise ValueError("Invalid background method specified.")
-
-        # Temporary file for processed image
+        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 = position_logic(temp_image_path, canvas_size, padding_top, padding_right, padding_bottom, padding_left)
 
-        # Create canvas based on background choice
+        # Create a new canvas with the appropriate background
         if bg_choice == 'white':
             canvas = Image.new("RGBA", canvas_size, "WHITE")
         elif bg_choice == 'custom':
@@ -241,62 +429,45 @@ def process_single_image(image_path, output_folder, bg_method, canvas_size_name,
         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"})
+        # 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"})
 
-        # 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
+        # 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.convert('RGB').save(output_path, format='JPEG')
+            canvas = canvas.convert('RGB')
+            canvas.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}")
+        print(f"Processed image path: {output_path}")
         return [(output_path, image_path)], log
 
     except Exception as e:
-        logging.error(f"Error processing {filename}: {e}")
+        print(f"Error processing {filename}: {e}")
         return None, None
-    
-# Set up logging
-logging.basicConfig(level=logging.INFO)
+
+def remove_extension(filename):
+    # Regular expression to match any extension at the end of the string
+    return re.sub(r'\.[^.]+$', '', filename)
 
 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)
@@ -305,8 +476,6 @@ def process_images(input_files, bg_method='rembg', watermark_path=None, canvas_s
     processed_images = []
     original_images = []
     all_logs = []
-    
-    image_files = []
 
     if isinstance(input_files, str) and input_files.lower().endswith(('.zip', '.rar')):
         # Handle zip file
@@ -318,19 +487,20 @@ def process_images(input_files, bg_method='rembg', watermark_path=None, canvas_s
         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}")
+            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)
-    logging.info(f"Total images to process: {total_images}")
+    print(f"Total images to process: {total_images}")
 
     avg_processing_time = 0
     with ThreadPoolExecutor(max_workers=num_workers) as executor:
@@ -341,22 +511,29 @@ def process_images(input_files, bg_method='rembg', watermark_path=None, canvas_s
                 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)
+                    if watermark_path:
+                        get_name = future_to_image[future].split('/')
+                        get_name = 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')
+                        add_twibbon(f'processed_images/{twibbon_input}', watermark_path, twibbon_output_path)
+                        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:
-                logging.error(f"Error processing image {future_to_image[future]}: {e}")
+                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:
@@ -366,40 +543,35 @@ def process_images(input_files, bg_method='rembg', watermark_path=None, canvas_s
     # 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'))
+    print("Comprehensive log saved to", 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")
-    
+    print(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
+
+    # 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 process_images(input_files, bg_method, watermark_path, canvas_size, output_format, bg_choice, custom_color, num_workers, progress)
+            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'))
+    if bg_choice == 'custom':
+        return gr.update(visible=True)
+    return gr.update(visible=False)
 
 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]
+        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)
@@ -408,25 +580,63 @@ def update_compare(evt: gr.SelectData):
         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)
-               )
+        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
+        print("No caption found in selection")
+        return gr.update(value=None), gr.update(value=None), gr.update(value=None), gr.update(value=None)
 
 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:
-        watermark_path = watermark.name if watermark else None        
-        _, 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"
+	_, 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"
+
+def add_twibbon(image_path, twibbon_path, output_path):
+    # Open the original image and the twibbon
+    image = Image.open(image_path)
+    twibbon = Image.open(twibbon_path)
+
+    # Get the sizes of both images
+    image_width, image_height = image.size
+    twibbon_width, twibbon_height = twibbon.size
+
+    # Resize the original image to fit inside the twibbon (optional: resize by aspect ratio)
+    aspect_ratio = image_width / image_height
+    if twibbon_width / twibbon_height > aspect_ratio:
+        new_width = twibbon_width
+        new_height = int(new_width / aspect_ratio)
+    else:
+        new_height = twibbon_height
+        new_width = int(new_height * aspect_ratio)
+
+    image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
+
+    # Center the image within the twibbon
+    x_offset = (twibbon_width - new_width) // 2
+    y_offset = (twibbon_height - new_height) // 2
+    combined_image = Image.new('RGBA', (twibbon_width, twibbon_height))
+    combined_image.paste(image, (x_offset, y_offset))
+    combined_image.paste(twibbon, (0, 0), mask=twibbon)  # Twibbon is pasted over the image
+
+    # Save the result
+    combined_image.save(output_path)
+    return combined_image
+
+def process_twibbon(image, twibbon):
+    output_path = "output_image.png"  # Output sementara
+    combined_image = add_twibbon(image.name, twibbon.name, output_path)
+    return combined_image
+
+def remove_background(image_path, method="none"):
+    image = Image.open(image_path)
+    
+    if method == "none":
+        return image  # Return the original image without any background removal
+    elif method == "rembg":
+        image = remove_background_rembg(image_path)
+    elif method == "bria":
+        image = remove_background_bria(image_path)
+    
+    return image  # Default return in case no valid method is chosen
 
 with gr.Blocks(theme="NoCrypt/miku@1.2.2") as iface:
     gr.Markdown("# 🎨 Creative Image Suite: Generate, Modify, and Enhance Your Visuals")
@@ -498,7 +708,7 @@ with gr.Blocks(theme="NoCrypt/miku@1.2.2") as iface:
         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", "none"], label="Background Removal Method", value="bria")
+        bg_method = gr.Radio(choices=["bria", "rembg", None], 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)