border color filter

app.py

@@ -9,9 +9,37 @@ import gradio as gr
 from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor
 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):
+    # Open the image
+    image = Image.open(image_path)
+    pixels = image.load()
+    
+    width, height = image.size
+
+    # Get the color of the first pixel on the left and right borders
+    left_border_color = pixels[0, 0]
+    right_border_color = pixels[width - 1, 0]
+
+    # Check the left border
+    for y in range(height):
+        if not colors_within_tolerance(pixels[0, y], left_border_color, tolerance):
+            return False
+    
+    # Check the right border
+    for y in range(height):
+        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}")
 
         # Calculate the scaling factor
         scaling_factor = max(target_size[0] / width, target_size[1] / height)
@@ -19,6 +47,7 @@ def resize_and_crop_image(image_path, target_size=(1080, 1080), crop_mode='cente
         # Resize the image with high-quality resampling
         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
@@ -41,10 +70,12 @@ def resize_and_crop_image(image_path, target_size=(1080, 1080), crop_mode='cente
 
         # Crop the image
         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)
@@ -52,28 +83,54 @@ def remove_background_rembg(input_path):
     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)  # applies mask on input and returns a pillow image
     return pillow_image
 
-def process_single_image(image_path, output_folder, crop_mode, remove_bg, bg_method, output_format, bg_choice, watermark_path=None):
+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:
-        if remove_bg == 'yes':
-            if bg_method == 'rembg':
-                # Remove background using rembg
-                image_with_no_bg = remove_background_rembg(image_path)
-            elif bg_method == 'bria':
-                # Remove background using 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')
+        print(f"Processing image: {filename}")
+        if bg_method == 'rembg':
+            # Remove background using rembg
+            image_with_no_bg = remove_background_rembg(image_path)
+        elif bg_method == 'bria':
+            # Remove background using 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')
+
+        # Check border colors and categorize
+        if check_border_colors(temp_image_path, tolerance=50):
+            print(f"Border colors are the same for image: {filename}")
+            # Create a new 1080x1080 canvas
+            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)
+
+            # Scale image to fit inside the canvas without stretching
+            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:
-            temp_image_path = image_path
-
-        # Resize and crop the image with or without background removal
-        new_image = resize_and_crop_image(temp_image_path, crop_mode=crop_mode)
+            print(f"Border colors are different for image: {filename}")
+            new_image = resize_and_crop_image(temp_image_path, crop_mode=crop_mode)
+
+        # Change background color if needed
+        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)
 
         # Save both versions of the image (with and without watermark)
         images_paths = []
@@ -99,17 +156,17 @@ def process_single_image(image_path, output_folder, crop_mode, remove_bg, bg_met
                 new_image_with_watermark.save(output_path_with_watermark, format='PNG')
             images_paths.append(output_path_with_watermark)
 
-        if remove_bg == 'yes':
-            # Remove the temporary file
-            os.remove(temp_image_path)
+        # Remove the temporary file
+        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', remove_bg='yes', bg_method='rembg', watermark_path=None, output_format='PNG', bg_choice='transparent', num_workers=4, progress=gr.Progress()):
+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()
     
     # Create a temporary directory
@@ -123,21 +180,29 @@ def process_images(zip_file, crop_mode='center', remove_bg='yes', bg_method='rem
     os.makedirs(output_folder)
     
     # Extract the zip file
-    with zipfile.ZipFile(zip_file, 'r') as zip_ref:
-        zip_ref.extractall(input_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 = []
     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}")
 
     # Process images using ThreadPoolExecutor for I/O-bound tasks and ProcessPoolExecutor for CPU-bound tasks
     with ThreadPoolExecutor(max_workers=num_workers) as thread_executor:
         with ProcessPoolExecutor(max_workers=num_workers) as process_executor:
-            future_to_image = {thread_executor.submit(process_single_image, image_path, output_folder, crop_mode, remove_bg, bg_method, output_format, bg_choice, watermark_path): image_path for image_path in image_files}
+            future_to_image = {thread_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):
-                result = future.result()
-                if result:
-                    processed_images.extend(result)
+                try:
+                    result = future.result()
+                    if result:
+                        processed_images.extend(result)
+                except Exception as e:
+                    print(f"Error processing image {future_to_image[future]}: {e}")
                 # Update progress
                 progress((idx + 1) / total_images, f"{idx + 1}/{total_images} images processed")
 
@@ -152,22 +217,23 @@ def process_images(zip_file, crop_mode='center', remove_bg='yes', bg_method='rem
 
     end_time = time.time()
     processing_time = end_time - start_time
+    print(f"Processing time: {processing_time} seconds")
     
     # Return the images, the zip file path, and processing time
     return processed_images, output_zip_path, processing_time
 
-def gradio_interface(zip_file, crop_mode, remove_bg, bg_method, watermark, output_format, bg_choice, num_workers):
+def gradio_interface(zip_file, crop_mode, bg_method, watermark, output_format, bg_choice, custom_color, num_workers):
     progress = gr.Progress()  # Initialize progress
     watermark_path = watermark.name if watermark else None
-    return process_images(zip_file.name, crop_mode, remove_bg, bg_method, watermark_path, output_format, bg_choice, num_workers, progress)
+    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(remove_bg, output_format):
-    if remove_bg == 'yes' and output_format == 'PNG':
+def show_bg_choice(output_format):
+    if output_format == 'PNG':
         return gr.update(visible=True)
     return gr.update(visible=False)
 
-def show_bg_method(remove_bg):
-    if remove_bg == 'yes':
+def show_color_picker(bg_choice):
+    if bg_choice == 'custom':
         return gr.update(visible=True)
     return gr.update(visible=False)
 
@@ -182,29 +248,28 @@ with gr.Blocks() as iface:
     
     with gr.Row():
         crop_mode = gr.Radio(choices=["center", "top", "bottom", "left", "right"], label="Crop Mode", value="center")
-        remove_bg = gr.Radio(choices=["yes", "no"], label="Remove Background", value="yes")
         
         output_format = gr.Radio(choices=["PNG", "JPG"], label="Output Format", value="PNG")
         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"], label="Background Choice", value="transparent", 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)
 
-    remove_bg.change(show_bg_choice, inputs=[remove_bg, output_format], outputs=bg_choice)
-    remove_bg.change(show_bg_method, inputs=remove_bg, outputs=bg_method)
-    output_format.change(show_bg_choice, inputs=[remove_bg, output_format], outputs=bg_choice)
+    bg_choice.change(show_color_picker, inputs=bg_choice, outputs=custom_color)
+    output_format.change(show_bg_choice, inputs=output_format, outputs=bg_choice)
 
     gallery = gr.Gallery(label="Processed Images")
     output_zip = gr.File(label="Download Processed Images as ZIP")
     processing_time = gr.Textbox(label="Processing Time (seconds)")
 
-    def process(zip_file, crop_mode, remove_bg, bg_method, watermark, output_format, bg_choice, num_workers):
-        processed_images, zip_path, time_taken = gradio_interface(zip_file, crop_mode, remove_bg, bg_method, watermark, output_format, bg_choice, num_workers)
+    def process(zip_file, crop_mode, bg_method, watermark, output_format, bg_choice, custom_color, num_workers):
+        processed_images, zip_path, time_taken = gradio_interface(zip_file, crop_mode, bg_method, watermark, output_format, bg_choice, custom_color, num_workers)
         return processed_images, zip_path, f"{time_taken:.2f} seconds"
 
     process_button = gr.Button("Process Images")
-    process_button.click(process, inputs=[zip_file, crop_mode, remove_bg, bg_method, watermark, output_format, bg_choice, num_workers], outputs=[gallery, output_zip, processing_time])
+    process_button.click(process, inputs=[zip_file, crop_mode, bg_method, watermark, output_format, bg_choice, custom_color, num_workers], outputs=[gallery, output_zip, processing_time])
 
 # Launch the interface
 iface.launch()

pixel_border.py

@@ -0,0 +1,50 @@
+import os
+import shutil
+from PIL import Image
+
+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):
+    # Open the image
+    image = Image.open(image_path)
+    pixels = image.load()
+    
+    width, height = image.size
+
+    # Get the color of the first pixel on the left and right borders
+    left_border_color = pixels[0, 0]
+    right_border_color = pixels[width - 1, 0]
+
+    # Check the left border
+    for y in range(height):
+        if not colors_within_tolerance(pixels[0, y], left_border_color, tolerance):
+            return False
+    
+    # Check the right border
+    for y in range(height):
+        if not colors_within_tolerance(pixels[width - 1, y], right_border_color, tolerance):
+            return False
+
+    return True
+
+def process_images(input_folder, output_folder_same, output_folder_different, tolerance):
+    # Ensure output directories exist
+    os.makedirs(output_folder_same, exist_ok=True)
+    os.makedirs(output_folder_different, exist_ok=True)
+
+    for filename in os.listdir(input_folder):
+        if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')):
+            image_path = os.path.join(input_folder, filename)
+            if check_border_colors(image_path, tolerance):
+                shutil.copy(image_path, os.path.join(output_folder_same, filename))
+            else:
+                shutil.copy(image_path, os.path.join(output_folder_different, filename))
+
+# Example usage
+input_folder = 'D:\Ardha\Kerja\AISensum\ROX\Task 1'
+output_folder_same = 'D:\Ardha\Kerja\AISensum\ROX\Task 1\warna_sama'
+output_folder_different = 'D:\Ardha\Kerja\AISensum\ROX\Task 1\warna_berbeda'
+tolerance = 50  # Adjust the tolerance value as needed
+
+process_images(input_folder, output_folder_same, output_folder_different, tolerance)