diff --git a/xiuren/random_pic.py b/xiuren/random_pic.py index cf1bd83..7613c78 100644 --- a/xiuren/random_pic.py +++ b/xiuren/random_pic.py @@ -5,15 +5,20 @@ import random def get_random_file_from_dir(directory): selected_file = None file_count = 0 + # 定义图片文件的扩展名列表 + image_extensions = {'.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp'} # 遍历目录及其所有子目录 for root, dirs, files in os.walk(directory): for file in files: - # 计算当前文件被选中的概率 - file_count += 1 - if random.randint(1, file_count) == 1: - # 每个文件有1/file_count的概率被选中 - selected_file = os.path.join(root, file) + # 获取文件扩展名 + _, ext = os.path.splitext(file) + # 只考虑图片文件 + if ext.lower() in image_extensions: + file_count += 1 + if random.randint(1, file_count) == 1: + # 每个图片文件有1/file_count的概率被选中 + selected_file = os.path.join(root, file) return selected_file @@ -22,4 +27,7 @@ def get_xiuren_pic(): # 使用示例 directory = '.' # 替换为你的目录路径 random_file_path = get_random_file_from_dir(directory) - return os.path.abspath(random_file_path) + if random_file_path: + return os.path.abspath(random_file_path) + else: + return None