From c64b58540de39a89c5d6a5e0fb16107c5f27ce5a Mon Sep 17 00:00:00 2001
From: liuwei <liuwei@wdtrgf.com.cn>
Date: Mon, 24 Feb 2025 16:36:40 +0800
Subject: [PATCH] =?UTF-8?q?=E9=80=89=E6=96=87=E4=BB=B6=E6=97=B6=EF=BC=8C?=
 =?UTF-8?q?=E5=8F=AA=E6=94=AF=E6=8C=81=E5=9B=BE=E7=89=87=E6=96=87=E4=BB=B6?=
 =?UTF-8?q?=EF=BC=8C=E9=98=B2=E6=AD=A2=E5=8F=91=E9=80=81=E5=88=AB=E7=9A=84?=
 =?UTF-8?q?=E8=B5=84=E6=96=99?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 xiuren/random_pic.py | 20 ++++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

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