import argparse import os import shutil from rembg import remove from PIL import Image import io def add_background(image, background, default_color="#FFFFFF"): """ Adds a background to an image, with a fallback to a default color if the specified background is not available. Args: - image (PIL.Image.Image): Image with a transparent background. - background (str or PIL.Image.Image): Background color (as a hex code) or a PIL Image to be used as background. - default_color (str): Fallback color if the specified background is not valid. Defaults to white. Returns: - PIL.Image.Image: The image with the new background. """ foreground = image.convert("RGBA") if isinstance(background, str) and (background.startswith("#") or background.isalpha()): # Background is a color try: Image.new("RGBA", (1, 1), background) # Test if valid color background_layer = Image.new("RGBA", foreground.size, background) except ValueError: print( f"Invalid color '{background}'. Using default color '{default_color}'.") background_layer = Image.new( "RGBA", foreground.size, default_color) elif isinstance(background, Image.Image): # Background is an image bg_img = background.convert("RGBA") background_layer = bg_img.resize(foreground.size) else: # Fallback to default color background_layer = Image.new("RGBA", foreground.size, default_color) final_img = Image.alpha_composite( background_layer, foreground).convert("RGB") return final_img def cropnontrans(image, padding=0): """ crops a nontransparent image Args: - image (PIL.Image.Image): Image to be cropped. Returns: - PIL.Image.Image: The autocropped image. """ # first, # Convert the PIL Image to bytes img_byte_arr = io.BytesIO() image.save(img_byte_arr, format='PNG') img_byte_arr = img_byte_arr.getvalue() # Use rembg to remove the background result_bytes = remove(img_byte_arr) # Convert the result bytes back to a PIL Image transparent_image = Image.open(io.BytesIO(result_bytes)) bbox = transparent_image.getbbox() # add padding area to the original bbox if bbox: bbox = (bbox[0]-padding, bbox[1]-padding, bbox[2]+padding, bbox[3]+padding) #delimig the bbox to the image size bbox = (max(0, bbox[0]), max(0, bbox[1]), min(transparent_image.width, bbox[2]), min(transparent_image.height, bbox[3])) print(f"Bounding box: {bbox}") return image.crop(bbox) return image def autocrop_image(image): """ Autocrops an image, focusing on the non-transparent pixels. Args: - image (PIL.Image.Image): Image to be autocropped. Returns: - PIL.Image.Image: The autocropped image. """ bbox = image.getbbox() print(f"Bounding box: {bbox}") if bbox: return image.crop(bbox) return image def remove_bg_func(image): """ Removes the background from an image using the rembg library. Args: - image (PIL.Image.Image): Image object from which to remove the background. Returns: - PIL.Image.Image: New image object with the background removed. """ # Convert the PIL Image to bytes img_byte_arr = io.BytesIO() image.save(img_byte_arr, format='PNG') img_byte_arr = img_byte_arr.getvalue() # Use rembg to remove the background result_bytes = remove(img_byte_arr) # Convert the result bytes back to a PIL Image result_image = Image.open(io.BytesIO(result_bytes)) return result_image # Update the process_image function to include format selection def process_image(img, crop=False, remove_bg=False, resize=None, padding=0, background=None, output_format='webp'): """ Processes a single image with the specified options and format. Args: - img: Input image - crop: Whether to autocrop - remove_bg: Whether to remove background - resize: Tuple of (width, height) for resizing - padding: Padding to add - background: Background color or image - output_format: 'webp', 'png', or 'png-transparent' """ # Existing processing code remains the same if remove_bg: img = remove_bg_func(img) if crop and remove_bg: img = autocrop_image(img) if crop and not remove_bg: img = cropnontrans(img, padding) if resize: img = resize_and_pad_image(img, resize, padding) if background and output_format != 'png-transparent': img = add_background(img, background) return img def resize_and_pad_image(image, dimensions, padding=0): """ Resizes an image to fit the specified dimensions and adds padding. Args: - image (PIL.Image.Image): Image object to be resized and padded. - dimensions (tuple): Target dimensions (width, height). - padding (int): Padding to add around the resized image. Returns: - PIL.Image.Image: Resized and padded image object. """ target_width, target_height = dimensions content_width, content_height = target_width - \ 2*padding, target_height - 2*padding # Determine new size, preserving aspect ratio img_ratio = image.width / image.height target_ratio = content_width / content_height if target_ratio > img_ratio: new_height = content_height new_width = int(new_height * img_ratio) else: new_width = content_width new_height = int(new_width / img_ratio) # Resize the image resized_img = image.resize( (new_width, new_height), Image.Resampling.LANCZOS) # Create a new image with the target dimensions and a transparent background new_img = Image.new( "RGBA", (target_width, target_height), (255, 255, 255, 0)) # Calculate the position to paste the resized image to center it paste_position = ((target_width - new_width) // 2, (target_height - new_height) // 2) # Paste the resized image onto the new image, centered new_img.paste(resized_img, paste_position, resized_img if resized_img.mode == 'RGBA' else None) return new_img def generate_output_filename(input_path, remove_bg=False, crop=False, resize=None, background=None): """ Generates an output filename based on the input path and processing options applied. Appends specific suffixes based on the operations: '_b' for background removal, '_c' for crop, and '_bg' if a background is added. It ensures the file extension is '.png'. Args: - input_path (str): Path to the input image. - remove_bg (bool): Indicates if background removal was applied. - crop (bool): Indicates if autocrop was applied. - resize (tuple): Optional dimensions (width, height) for resizing the image. - background (str): Indicates if a background was added (None if not used). Returns: - (str): Modified filename with appropriate suffix and '.png' extension. """ base, _ = os.path.splitext(os.path.basename(input_path)) suffix = "" if remove_bg: suffix += "_b" if crop: suffix += "_c" if resize: width, height = resize suffix += f"_{width}x{height}" if background: suffix += "_bg" # Append "_bg" if the background option was used # Ensure the file saves as PNG, accommodating for transparency or added backgrounds return f"{base}{suffix}.png" # The main and process_images functions remain the same, but ensure to update them to handle the new PNG output correctly. # Update the process_images and main functions to include the new autocrop functionality # Ensure to pass the crop argument to process_image and adjust the output filename generation accordingly def process_images(input_dir="./input", output_dir="./output", crop=False, remove_bg=False, resize=None, padding=0, background=None): """ Processes images in the specified directory based on the provided options. """ processed_input_dir = os.path.join(input_dir, "processed") os.makedirs(processed_input_dir, exist_ok=True) os.makedirs(output_dir, exist_ok=True) inputs = [os.path.join(input_dir, f) for f in os.listdir( input_dir) if os.path.isfile(os.path.join(input_dir, f))] if not inputs: print("No images found in the input directory.") return for i, input_path in enumerate(inputs, start=1): try: with Image.open(input_path) as img: # Define filename here, before it's used filename = os.path.basename(input_path) # Process the image processed_img = process_image( img, crop=crop, remove_bg=remove_bg, resize=resize, padding=padding, background=background) # Generate output filename based on processing parameters output_filename = generate_output_filename( filename, remove_bg=remove_bg, crop=crop, resize=resize, background=background) output_path = os.path.join(output_dir, output_filename) # Save the processed image to the output directory processed_img.save(output_path) print( f"Processed image {i}/{len(inputs)}: {filename} -> {output_filename}") # Optionally move the processed input image to a "processed" subdirectory shutil.move(input_path, os.path.join( processed_input_dir, filename)) except Exception as e: print(f"Error processing image {input_path}: {e}") print("All images have been processed.") def save_image_with_format(image, output_path, format='webp', quality=90, custom_filename=None): """ Saves the image in the specified format with appropriate settings. Args: - image (PIL.Image.Image): The image to save - output_path (str): Base path for the output file (without extension) - format (str): 'webp', 'png', 'png-transparent', or 'jpg' - quality (int): Quality setting for compression (1-100) - custom_filename (str): Optional custom filename for the output """ # Get image dimensions for filename width, height = image.size # Generate filename with schema: originalname_size.type if custom_filename: base_dir = os.path.dirname(output_path) filename = f"{custom_filename}_{width}x{height}" final_path = os.path.join(base_dir, filename) else: final_path = output_path if format == 'webp': final_path = f"{final_path}.webp" image.save(final_path, 'webp', quality=quality) elif format == 'png-transparent': final_path = f"{final_path}.png" image.save(final_path, 'PNG', optimize=True) elif format == 'png': final_path = f"{final_path}.png" if image.mode in ('RGBA', 'LA'): background = Image.new('RGB', image.size, 'white') background.paste(image, mask=image.split()[-1]) background.save(final_path, 'PNG', optimize=True) else: image.save(final_path, 'PNG', optimize=True) elif format == 'jpg': final_path = f"{final_path}.jpg" if image.mode in ('RGBA', 'LA'): background = Image.new('RGB', image.size, 'white') background.paste(image, mask=image.split()[-1]) background.save(final_path, 'JPEG', quality=quality, optimize=True) else: image.convert('RGB').save(final_path, 'JPEG', quality=quality, optimize=True) else: raise ValueError(f"Unsupported format: {format}") return final_path