add thread

app.py

@@ -5,6 +5,7 @@ from PIL import Image
 import io
 from rembg import remove
 import gradio as gr
+from concurrent.futures import ThreadPoolExecutor
 
 def resize_and_crop_image(image_path, target_size=(1080, 1080), crop_mode='center'):
     with Image.open(image_path) as img:
@@ -48,6 +49,34 @@ def remove_background(input_path):
     img = Image.open(io.BytesIO(output_image)).convert("RGBA")
     return img
 
+def process_single_image(image_path, output_folder, crop_mode, remove_bg):
+    filename = os.path.basename(image_path)
+    try:
+        if remove_bg == 'yes':
+            # Remove background
+            image_with_no_bg = remove_background(image_path)
+            temp_image_path = os.path.join(output_folder, f"temp_{filename}")
+            image_with_no_bg.save(temp_image_path, format='PNG')
+        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)
+
+        # Save the final image
+        output_path = os.path.join(output_folder, f"{os.path.splitext(filename)[0]}.png")
+        new_image.save(output_path, format='PNG')
+
+        if remove_bg == 'yes':
+            # Remove the temporary file
+            os.remove(temp_image_path)
+
+        return output_path
+
+    except Exception as e:
+        print(f"Error processing {filename}: {e}")
+        return None
+
 def process_images(zip_file, crop_mode='center', remove_bg='yes', progress=gr.Progress()):
     # Create a temporary directory
     input_folder = "temp_input"
@@ -64,37 +93,18 @@ def process_images(zip_file, crop_mode='center', remove_bg='yes', progress=gr.Pr
         zip_ref.extractall(input_folder)
     
     processed_images = []
-    image_files = [f for f in os.listdir(input_folder) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif'))]
+    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)
-    for idx, filename in enumerate(image_files):
-        image_path = os.path.join(input_folder, filename)
-        try:
-            if remove_bg == 'yes':
-                # Remove background
-                image_with_no_bg = remove_background(image_path)
-                temp_image_path = os.path.join(output_folder, f"temp_{filename}")
-                image_with_no_bg.save(temp_image_path, format='PNG')
-            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)
-
-            # Save the final image
-            output_path = os.path.join(output_folder, f"{os.path.splitext(filename)[0]}.png")
-            new_image.save(output_path, format='PNG')
-
-            processed_images.append(output_path)
-
-            if remove_bg == 'yes':
-                # Remove the temporary file
-                os.remove(temp_image_path)
 
+    # 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): image_path for image_path in image_files}
+        for idx, future in enumerate(future_to_image):
+            result = future.result()
+            if result:
+                processed_images.append(result)
             # Update progress
             progress((idx + 1) / total_images, f"{idx + 1}/{total_images} images processed")
-
-        except Exception as e:
-            print(f"Error processing {filename}: {e}")
     
     # Create a zip file of the processed images
     output_zip_path = "processed_images.zip"