change thread logic

app.py

@@ -1,15 +1,15 @@
 import os
 import zipfile
 import shutil
+import time
 from PIL import Image
 import io
 from rembg import remove
 import gradio as gr
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor
 from transformers import pipeline
 
 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
 
@@ -45,7 +45,6 @@ def resize_and_crop_image(image_path, target_size=(1080, 1080), crop_mode='cente
         return cropped_img
 
 def remove_background_rembg(input_path):
-    print(f"Removing background using rembg: {input_path}")
     with open(input_path, 'rb') as i:
         input_image = i.read()
     output_image = remove(input_image)
@@ -53,7 +52,6 @@ def remove_background_rembg(input_path):
     return img
 
 def remove_background_bria(input_path):
-    print(f"Removing background using bria: {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
@@ -61,7 +59,6 @@ def remove_background_bria(input_path):
 def process_single_image(image_path, output_folder, crop_mode, remove_bg, bg_method, output_format, bg_choice, watermark_path=None):
     filename = os.path.basename(image_path)
     try:
-        print(f"Processing image: {filename}")
         if remove_bg == 'yes':
             if bg_method == 'rembg':
                 # Remove background using rembg
@@ -72,13 +69,11 @@ def process_single_image(image_path, output_folder, crop_mode, remove_bg, bg_met
                 
             temp_image_path = os.path.join(output_folder, f"temp_{filename}")
             image_with_no_bg.save(temp_image_path, format='PNG')
-            print(f"Background removed and saved temporary image: {temp_image_path}")
         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"Resized and cropped image: {temp_image_path}")
 
         # Save both versions of the image (with and without watermark)
         images_paths = []
@@ -91,7 +86,6 @@ def process_single_image(image_path, output_folder, crop_mode, remove_bg, bg_met
         else:
             new_image.save(output_path_without_watermark, format='PNG')
         images_paths.append(output_path_without_watermark)
-        print(f"Saved final image without watermark: {output_path_without_watermark}")
 
         # Apply watermark if provided and save the version with watermark
         if watermark_path:
@@ -104,12 +98,10 @@ def process_single_image(image_path, output_folder, crop_mode, remove_bg, bg_met
             else:
                 new_image_with_watermark.save(output_path_with_watermark, format='PNG')
             images_paths.append(output_path_with_watermark)
-            print(f"Saved final image with watermark: {output_path_with_watermark}")
 
         if remove_bg == 'yes':
             # Remove the temporary file
             os.remove(temp_image_path)
-            print(f"Removed temporary file: {temp_image_path}")
 
         return images_paths
 
@@ -117,8 +109,9 @@ def process_single_image(image_path, output_folder, crop_mode, remove_bg, bg_met
         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', progress=gr.Progress()):
-    print("Starting to process images...")
+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()):
+    start_time = time.time()
+    
     # Create a temporary directory
     input_folder = "temp_input"
     output_folder = "temp_output"
@@ -130,25 +123,23 @@ def process_images(zip_file, crop_mode='center', remove_bg='yes', bg_method='rem
     os.makedirs(output_folder)
     
     # Extract the zip file
-    print(f"Extracting zip file: {zip_file}")
     with zipfile.ZipFile(zip_file, 'r') as zip_ref:
         zip_ref.extractall(input_folder)
     
     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"Found {total_images} images to process.")
-
-    # Process images using ThreadPoolExecutor
-    with ThreadPoolExecutor(max_workers=2) as executor:
-        future_to_image = {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}
-        for idx, future in enumerate(future_to_image):
-            result = future.result()
-            if result:
-                processed_images.extend(result)
-            # Update progress
-            progress((idx + 1) / total_images, f"{idx + 1}/{total_images} images processed")
-            print(f"Processed {idx + 1}/{total_images} 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}
+            for idx, future in enumerate(future_to_image):
+                result = future.result()
+                if result:
+                    processed_images.extend(result)
+                # Update progress
+                progress((idx + 1) / total_images, f"{idx + 1}/{total_images} images processed")
 
     # Create a zip file of the processed images
     output_zip_path = "processed_images.zip"
@@ -159,15 +150,16 @@ def process_images(zip_file, crop_mode='center', remove_bg='yes', bg_method='rem
             else:
                 zipf.write(file, os.path.join("without_watermark", os.path.basename(file)))
 
-    print(f"Created zip file of processed images: {output_zip_path}")
-    # Return the images and the zip file path
-    return processed_images, output_zip_path
+    end_time = time.time()
+    processing_time = end_time - start_time
+    
+    # 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):
-    print("Initializing Gradio interface processing...")
+def gradio_interface(zip_file, crop_mode, remove_bg, bg_method, watermark, output_format, bg_choice, 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, progress)
+    return process_images(zip_file.name, crop_mode, remove_bg, bg_method, watermark_path, output_format, bg_choice, num_workers, progress)
 
 def show_bg_choice(remove_bg, output_format):
     if remove_bg == 'yes' and output_format == 'PNG':
@@ -192,7 +184,8 @@ with gr.Blocks() as iface:
         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")    
+        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=4)
 
     with gr.Row():
         bg_method = gr.Radio(choices=["bria", "rembg"], label="Background Removal Method", value="bria", visible=True)
@@ -204,13 +197,14 @@ with gr.Blocks() as iface:
 
     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):
-        processed_images, zip_path = gradio_interface(zip_file, crop_mode, remove_bg, bg_method, watermark, output_format, bg_choice)
-        return processed_images, zip_path
+    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)
+        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], outputs=[gallery, output_zip])
+    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])
 
 # Launch the interface
 iface.launch()