Spaces:
Sleeping
Sleeping
| import os | |
| import zipfile | |
| import shutil | |
| from PIL import Image | |
| import io | |
| from rembg import remove | |
| import gradio as gr | |
| 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_images(zip_file, crop_mode='center', remove_bg='yes', 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 = [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) | |
| # 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" | |
| 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): | |
| progress = gr.Progress() # Initialize progress | |
| return process_images(zip_file.name, crop_mode, remove_bg, progress) | |
| # Create the Gradio interface | |
| with gr.Blocks() as iface: | |
| gr.Markdown("# Image Background Removal and Resizing") | |
| gr.Markdown("Upload a ZIP or RAR file containing images, choose the crop mode, and get the processed images.") | |
| with gr.Row(): | |
| zip_file = gr.File(label="Upload ZIP/RAR file of images", file_types=[".zip", ".rar"]) | |
| 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") | |
| gallery = gr.Gallery(label="Processed Images") | |
| output_zip = gr.File(label="Download Processed Images as ZIP") | |
| def process(zip_file, crop_mode, remove_bg): | |
| processed_images, zip_path = gradio_interface(zip_file, crop_mode, remove_bg) | |
| return processed_images, zip_path | |
| process_button = gr.Button("Process Images") | |
| process_button.click(process, inputs=[zip_file, crop_mode, remove_bg], outputs=[gallery, output_zip]) | |
| # Launch the interface | |
| iface.launch() | |