Spaces:
Sleeping
Sleeping
| 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 resize_and_crop_image(image_path, target_size=(1080, 1080), crop_mode='center'): | |
| with Image.open(image_path) as img: | |
| width, height = img.size | |
| # 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) | |
| 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)) | |
| return cropped_img | |
| def remove_background_rembg(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): | |
| 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, remove_bg, bg_method, output_format, bg_choice, watermark_path=None): | |
| filename = os.path.basename(image_path) | |
| try: | |
| if remove_bg == 'yes': | |
| 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') | |
| 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 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) | |
| if remove_bg == 'yes': | |
| # Remove the temporary file | |
| os.remove(temp_image_path) | |
| return images_paths | |
| except Exception as e: | |
| 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', 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 | |
| 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) | |
| # 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" | |
| 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 | |
| # 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, 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, num_workers, progress) | |
| def show_bg_choice(remove_bg, output_format): | |
| if remove_bg == 'yes' and output_format == 'PNG': | |
| return gr.update(visible=True) | |
| return gr.update(visible=False) | |
| def show_bg_method(remove_bg): | |
| if remove_bg == 'yes': | |
| 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") | |
| remove_bg = gr.Radio(choices=["yes", "no"], label="Remove Background", value="yes") | |
| 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"], label="Background Choice", value="transparent", visible=True) | |
| remove_bg.change(show_bg_choice, inputs=[remove_bg, output_format], outputs=bg_choice) | |
| remove_bg.change(show_bg_method, inputs=remove_bg, outputs=bg_method) | |
| output_format.change(show_bg_choice, inputs=[remove_bg, 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, 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, num_workers], outputs=[gallery, output_zip, processing_time]) | |
| # Launch the interface | |
| iface.launch() | |