change tread logic

app.py

@@ -6,30 +6,24 @@ from PIL import Image
 import io
 from rembg import remove
 import gradio as gr
-from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor
+from concurrent.futures import ThreadPoolExecutor
 from transformers import pipeline
 
 def colors_within_tolerance(color1, color2, tolerance):
     return all(abs(c1 - c2) <= tolerance for c1, c2 in zip(color1, color2))
 
 def check_border_colors(image_path, tolerance):
-    # 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
 
@@ -41,10 +35,8 @@ def resize_and_crop_image(image_path, target_size=(1080, 1080), crop_mode='cente
         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)
 
-        # 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}")
@@ -68,7 +60,6 @@ def resize_and_crop_image(image_path, target_size=(1080, 1080), crop_mode='cente
         right = left + target_size[0]
         bottom = top + target_size[1]
 
-        # Crop the image
         cropped_img = resized_img.crop((left, top, right, bottom))
         print(f"Cropped image size: {cropped_img.size}")
 
@@ -85,7 +76,7 @@ def remove_background_rembg(input_path):
 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
+    pillow_image = pipe(input_path)
     return pillow_image
 
 def process_single_image(image_path, output_folder, crop_mode, bg_method, output_format, bg_choice, custom_color, watermark_path=None):
@@ -93,25 +84,20 @@ def process_single_image(image_path, output_folder, crop_mode, bg_method, output
     try:
         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))
@@ -122,7 +108,6 @@ def process_single_image(image_path, output_folder, crop_mode, bg_method, output
             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")
@@ -132,10 +117,8 @@ def process_single_image(image_path, output_folder, crop_mode, bg_method, output
             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 = []
 
-        # Save without watermark
         output_ext = 'jpg' if output_format == 'JPG' else 'png'
         output_path_without_watermark = os.path.join(output_folder, f"without_watermark_{os.path.splitext(filename)[0]}.{output_ext}")
         if output_format == 'JPG':
@@ -144,7 +127,6 @@ def process_single_image(image_path, output_folder, crop_mode, bg_method, output
             new_image.save(output_path_without_watermark, format='PNG')
         images_paths.append(output_path_without_watermark)
 
-        # Apply watermark if provided and save the version with watermark
         if watermark_path:
             watermark = Image.open(watermark_path).convert("RGBA")
             new_image_with_watermark = new_image.copy()
@@ -156,7 +138,6 @@ def process_single_image(image_path, output_folder, crop_mode, bg_method, output
                 new_image_with_watermark.save(output_path_with_watermark, format='PNG')
             images_paths.append(output_path_with_watermark)
 
-        # Remove the temporary file
         os.remove(temp_image_path)
 
         print(f"Processed image paths: {images_paths}")
@@ -169,7 +150,6 @@ def process_single_image(image_path, output_folder, crop_mode, bg_method, output
 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
     input_folder = "temp_input"
     output_folder = "temp_output"
     if os.path.exists(input_folder):
@@ -179,7 +159,6 @@ def process_images(zip_file, crop_mode='center', bg_method='rembg', watermark_pa
     os.makedirs(input_folder)
     os.makedirs(output_folder)
     
-    # Extract the zip file
     try:
         with zipfile.ZipFile(zip_file, 'r') as zip_ref:
             zip_ref.extractall(input_folder)
@@ -188,25 +167,23 @@ def process_images(zip_file, crop_mode='center', bg_method='rembg', watermark_pa
         return [], None, 0
     
     processed_images = []
+    original_images = []
     image_files = [os.path.join(input_folder, f) for f in os.listdir(input_folder) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif'))]
     total_images = len(image_files)
     print(f"Total images to process: {total_images}")
 
-    # 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, bg_method, output_format, bg_choice, custom_color, watermark_path): image_path for image_path in image_files}
-            for idx, future in enumerate(future_to_image):
-                try:
-                    result = future.result()
-                    if result:
-                        processed_images.extend(result)
-                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")
-
-    # Create a zip file of the processed images
+    with ThreadPoolExecutor(max_workers=num_workers) as executor:
+        future_to_image = {executor.submit(process_single_image, image_path, output_folder, crop_mode, bg_method, output_format, bg_choice, custom_color, watermark_path): image_path for image_path in image_files}
+        for idx, future in enumerate(future_to_image):
+            try:
+                result = future.result()
+                if result:
+                    processed_images.extend(result)
+                    original_images.append(future_to_image[future])
+            except Exception as e:
+                print(f"Error processing image {future_to_image[future]}: {e}")
+            progress((idx + 1) / total_images, f"{idx + 1}/{total_images} images processed")
+
     output_zip_path = "processed_images.zip"
     with zipfile.ZipFile(output_zip_path, 'w') as zipf:
         for file in processed_images:
@@ -219,11 +196,10 @@ def process_images(zip_file, crop_mode='center', bg_method='rembg', watermark_pa
     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
+    return original_images, processed_images, output_zip_path, processing_time
 
 def gradio_interface(zip_file, crop_mode, bg_method, watermark, output_format, bg_choice, custom_color, num_workers):
-    progress = gr.Progress()  # Initialize progress
+    progress = gr.Progress()
     watermark_path = watermark.name if watermark else None
     return process_images(zip_file.name, crop_mode, bg_method, watermark_path, output_format, bg_choice, custom_color, num_workers, progress)
 
@@ -237,7 +213,6 @@ def show_color_picker(bg_choice):
         return gr.update(visible=True)
     return gr.update(visible=False)
 
-# Create the Gradio interface
 with gr.Blocks() as iface:
     gr.Markdown("# Image Background Removal and Resizing with Optional Watermark")
     gr.Markdown("Upload a ZIP or RAR file containing images, choose the crop mode, optionally upload a watermark image, and select the output format.")
@@ -248,7 +223,6 @@ with gr.Blocks() as iface:
     
     with gr.Row():
         crop_mode = gr.Radio(choices=["center", "top", "bottom", "left", "right"], label="Crop Mode", value="center")
-        
         output_format = gr.Radio(choices=["PNG", "JPG"], label="Output Format", value="PNG")
         num_workers = gr.Slider(minimum=1, maximum=16, step=1, label="Number of Workers", value=2)
 
@@ -257,19 +231,21 @@ with gr.Blocks() as iface:
         bg_choice = gr.Radio(choices=["transparent", "white", "custom"], label="Background Choice", value="transparent", visible=True)
         custom_color = gr.ColorPicker(label="Custom Background Color", value="#ffffff", visible=False)
 
+    with gr.Row():
+        gallery_original = gr.Gallery(label="Original Images")
+        gallery_processed = gr.Gallery(label="Processed Images")
+    with gr.Row():
+        output_zip = gr.File(label="Download Processed Images as ZIP")
+        processing_time = gr.Textbox(label="Processing Time (seconds)")
+
     bg_choice.change(show_color_picker, inputs=bg_choice, outputs=custom_color)
     output_format.change(show_bg_choice, inputs=output_format, outputs=bg_choice)
 
-    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, 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"
+        original_images, processed_images, zip_path, time_taken = gradio_interface(zip_file, crop_mode, bg_method, watermark, output_format, bg_choice, custom_color, num_workers)
+        return original_images, processed_images, zip_path, f"{time_taken:.2f} seconds"
 
     process_button = gr.Button("Process Images")
-    process_button.click(process, inputs=[zip_file, crop_mode, bg_method, watermark, output_format, bg_choice, custom_color, num_workers], outputs=[gallery, 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_original, gallery_processed, output_zip, processing_time])
 
-# Launch the interface
 iface.launch()