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, 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) # 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 top = (resized_img.height - target_size[1]) / 2 elif crop_mode == 'top': left = (resized_img.width - target_size[0]) / 2 top = 0 elif crop_mode == 'bottom': left = (resized_img.width - target_size[0]) / 2 top = resized_img.height - target_size[1] elif crop_mode == 'left': left = 0 top = (resized_img.height - target_size[1]) / 2 elif crop_mode == 'right': left = resized_img.width - target_size[0] top = (resized_img.height - target_size[1]) / 2 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}") 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) img = Image.open(io.BytesIO(output_image)).convert("RGBA") 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, bg_method, output_format, bg_choice, custom_color, watermark_path=None): filename = os.path.basename(image_path) 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)) 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: 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 = [] # 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': new_image.convert('RGB').save(output_path_without_watermark, format='JPEG') else: 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() new_image_with_watermark.paste(watermark, (0, 0), watermark) output_path_with_watermark = os.path.join(output_folder, f"with_watermark_{os.path.splitext(filename)[0]}.{output_ext}") if output_format == 'JPG': new_image_with_watermark.convert('RGB').save(output_path_with_watermark, format='JPEG') else: 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}") return images_paths except Exception as e: print(f"Error processing {filename}: {e}") return None 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): shutil.rmtree(input_folder) if os.path.exists(output_folder): shutil.rmtree(output_folder) 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) 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, 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 output_zip_path = "processed_images.zip" with zipfile.ZipFile(output_zip_path, 'w') as zipf: for file in processed_images: if "with_watermark" in file: zipf.write(file, os.path.join("with_watermark", os.path.basename(file))) else: zipf.write(file, os.path.join("without_watermark", os.path.basename(file))) 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, 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, bg_method, watermark_path, output_format, bg_choice, custom_color, num_workers, progress) def show_bg_choice(output_format): if output_format == 'PNG': return gr.update(visible=True) return gr.update(visible=False) def show_color_picker(bg_choice): if bg_choice == 'custom': 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.") with gr.Row(): zip_file = gr.File(label="Upload ZIP/RAR file of images", file_types=[".zip", ".rar"]) watermark = gr.File(label="Upload Watermark Image (Optional)", file_types=[".png"]) 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) 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", "custom"], label="Background Choice", value="transparent", visible=True) custom_color = gr.ColorPicker(label="Custom Background Color", value="#ffffff", visible=False) 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" 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]) # Launch the interface iface.launch()