Spaces:
Sleeping
Sleeping
| import os | |
| import zipfile | |
| import shutil | |
| 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: | |
| 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(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 process_single_image(image_path, output_folder, crop_mode, remove_bg, output_format, bg_choice, watermark_path=None): | |
| 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) | |
| # Apply watermark if provided | |
| if watermark_path: | |
| watermark = Image.open(watermark_path).convert("RGBA") | |
| new_image.paste(watermark, (0, 0), watermark) | |
| # Save the final image | |
| if remove_bg == 'yes' and output_format == 'PNG' and bg_choice == 'transparent': | |
| output_path = os.path.join(output_folder, f"{os.path.splitext(filename)[0]}.png") | |
| new_image.save(output_path, format='PNG') | |
| elif remove_bg == 'yes' and output_format == 'PNG' and bg_choice == 'white': | |
| # Create a white background | |
| background = Image.new("RGBA", new_image.size, (255, 255, 255, 255)) | |
| # Paste the image with no background onto the white background | |
| background.paste(new_image, mask=new_image.split()[3]) # 3 is the alpha channel | |
| output_path = os.path.join(output_folder, f"{os.path.splitext(filename)[0]}.png") | |
| background.convert('RGB').save(output_path, format='PNG') | |
| elif remove_bg == 'yes' and output_format == 'JPG': | |
| # Create a white background | |
| background = Image.new("RGB", new_image.size, (255, 255, 255)) | |
| # Paste the image with no background onto the white background | |
| background.paste(new_image, mask=new_image.split()[3]) # 3 is the alpha channel | |
| output_path = os.path.join(output_folder, f"{os.path.splitext(filename)[0]}.jpg") | |
| background.save(output_path, format='JPEG') | |
| else: | |
| output_ext = 'jpg' if output_format == 'JPG' else 'png' | |
| output_path = os.path.join(output_folder, f"{os.path.splitext(filename)[0]}.{output_ext}") | |
| if output_format == 'JPG': | |
| new_image.convert('RGB').save(output_path, format='JPEG') | |
| else: | |
| 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', watermark_path=None, output_format='PNG', bg_choice='transparent', progress=gr.Progress()): | |
| # 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 | |
| with ThreadPoolExecutor(max_workers=2) as executor: | |
| future_to_image = {executor.submit(process_single_image, image_path, output_folder, crop_mode, remove_bg, 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.append(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: | |
| zipf.write(file, os.path.basename(file)) | |
| # Return the images and the zip file path | |
| return processed_images, output_zip_path | |
| def gradio_interface(zip_file, crop_mode, remove_bg, watermark, output_format, bg_choice): | |
| progress = gr.Progress() # Initialize progress | |
| watermark_path = watermark.name if watermark else None | |
| return process_images(zip_file.name, crop_mode, remove_bg, watermark_path, output_format, bg_choice, 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) | |
| # 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") | |
| 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) | |
| 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") | |
| def process(zip_file, crop_mode, remove_bg, watermark, output_format, bg_choice): | |
| processed_images, zip_path = gradio_interface(zip_file, crop_mode, remove_bg, watermark, output_format, bg_choice) | |
| return processed_images, zip_path | |
| process_button = gr.Button("Process Images") | |
| process_button.click(process, inputs=[zip_file, crop_mode, remove_bg, watermark, output_format, bg_choice], outputs=[gallery, output_zip]) | |
| # Launch the interface | |
| iface.launch() | |