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 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): image = Image.open(image_path) pixels = image.load() width, height = image.size left_border_color = pixels[0, 0] right_border_color = pixels[width - 1, 0] for y in range(height): if not colors_within_tolerance(pixels[0, y], left_border_color, tolerance): return False 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}") scaling_factor = max(target_size[0] / width, target_size[1] / height) 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] 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) 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': image_with_no_bg = remove_background_rembg(image_path) elif bg_method == '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') if check_border_colors(temp_image_path, tolerance=50): print(f"Border colors are the same for image: {filename}") 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) 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) 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) images_paths = [] 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) 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) 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() 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) 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 = [] 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}") 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: 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 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() 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) 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) 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) def process(zip_file, crop_mode, bg_method, watermark, output_format, bg_choice, custom_color, num_workers): 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_original, gallery_processed, output_zip, processing_time]) iface.launch()