utils/count.py

import os

def count_images(directory):
    # 支持的图片文件扩展名
    image_extensions = {'.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp'}
    image_count = 0
    
    # 遍历指定目录及其子目录
    for root, dirs, files in os.walk(directory):
        for file in files:
            # 检查文件扩展名是否在支持的扩展名集合中
            if os.path.splitext(file)[1].lower() in image_extensions:
                image_count += 1
                
    return image_count

# 示例用法：替换 'path_to_your_directory' 为你的目录路径
directory_path = '/mnt/petrelfs/zhuchenglin/diffusion/images_large'
print(f"Total images in '{directory_path}':", count_images(directory_path))