秀人网图片随机优化

xiuren/random_pic.py

@@ -3,24 +3,17 @@ import random
 
 
 def get_random_file_from_dir(directory):
-    selected_file = None
-    file_count = 0
-    # 定义图片文件的扩展名列表
     image_extensions = {'.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp'}
+    image_files = []
 
-    # 遍历目录及其所有子目录
-    for root, dirs, files in os.walk(directory):
-        for file in files:
-            # 获取文件扩展名
-            _, 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)
+    with os.scandir(directory) as entries:
+        for entry in entries:
+            if entry.is_file():
+                _, ext = os.path.splitext(entry.name)
+                if ext.lower() in image_extensions:
+                    image_files.append(entry.path)
 
-    return selected_file
+    return random.choice(image_files) if image_files else None
 
 
 def get_xiuren_pic():