mirror of
https://github.com/karl0ss/comfy_fm24_newgens.git
synced 2025-04-29 04:13:40 +01:00
76 lines
2.6 KiB
Python
76 lines
2.6 KiB
Python
"""
|
|
Author: Emil Mirzayev
|
|
|
|
This script removes the background from images in a specified directory and its subfolders.
|
|
It processes JPG and JPEG files, converting them to PNG files with transparent backgrounds.
|
|
"""
|
|
|
|
import os
|
|
import argparse
|
|
from rembg import remove
|
|
from PIL import Image
|
|
import glob
|
|
from tqdm import tqdm
|
|
|
|
|
|
|
|
def remove_background(input_path, output_path):
|
|
"""
|
|
Remove the background from a single image.
|
|
|
|
Args:
|
|
input_path (str): Path to the input image file.
|
|
output_path (str): Path where the processed image will be saved.
|
|
"""
|
|
with Image.open(input_path) as img:
|
|
output = remove(img)
|
|
output.save(output_path)
|
|
|
|
def process_directory(directory):
|
|
"""
|
|
Process all JPG and JPEG images in the given directory and its subfolders.
|
|
|
|
Args:
|
|
directory (str): Path to the directory containing images.
|
|
|
|
Returns:
|
|
int: The number of images successfully processed.
|
|
"""
|
|
processed_count = 0
|
|
|
|
# Get the total number of files to process
|
|
total_files = sum(len(files) for _, _, files in os.walk(directory))
|
|
|
|
# Create a progress bar
|
|
with tqdm(total=total_files, desc="Processing images", unit="image") as pbar:
|
|
for subdir, dirs, files in os.walk(directory):
|
|
for file in files:
|
|
if file.lower().endswith(('.jpg', '.jpeg', 'png')):
|
|
input_path = os.path.join(subdir, file)
|
|
output_filename = os.path.splitext(file)[0] + '.png'
|
|
output_path = os.path.join(subdir, output_filename)
|
|
|
|
try:
|
|
remove_background(input_path, output_path)
|
|
processed_count += 1
|
|
except Exception as e:
|
|
print(f"Error processing {input_path}: {str(e)}")
|
|
|
|
# Update the progress bar
|
|
pbar.update(1)
|
|
|
|
return processed_count
|
|
def main():
|
|
"""
|
|
Main function to parse command-line arguments and initiate the background removal process.
|
|
"""
|
|
parser = argparse.ArgumentParser(description="Remove background from images in a directory and its subfolders")
|
|
parser.add_argument("--directory", type=str, default="generated_images", help="Path to the directory containing images. Defaults to `generated_images` folder in the same directory")
|
|
args = parser.parse_args()
|
|
|
|
total_processed = process_directory(args.directory)
|
|
print(f"\nTotal images processed: {total_processed}")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|