mirror of
https://github.com/karl0ss/comfy_fm24_newgens.git
synced 2025-04-29 04:13:40 +01:00
41 lines
1.4 KiB
Python
41 lines
1.4 KiB
Python
|
import os
|
||
|
|
from rembg import remove
|
||
|
|
from PIL import Image
|
||
|
|
from tqdm import tqdm
|
||
|
|
|
||
|
|
def remove_bg_from_files_in_dir(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:
|
||
|
|
with Image.open(input_path) as img:
|
||
|
|
output = remove(img)
|
||
|
|
output.save(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
|