diff --git a/function.py b/function.py
index 54be413..de3548a 100644
--- a/function.py
+++ b/function.py
@@ -30,12 +30,12 @@ def classify(input_vector, dataSet, labels, k=3):
     # 返回投票最多的标签
     return max(label_count, key=label_count.get)
 
-def load_data(directory):
+def load_dataset(directory):
     """
     加载手写数字数据集.
 
     Args:
-    - directory: D:\二值手写数字数据集D:\二值手写数字数据集
+    - directory: 数据集所在目录.
 
     Returns:
     - 特征数据集和标签集.
@@ -72,13 +72,16 @@ def img_to_vector(filename):
                 vector[0, 32 * i + j] = int(line[j])
     return vector
 
-
 if __name__ == "__main__":
+    # 设置数据集目录路径
+    train_data_dir = r"D:\二值手写数字数据集\trainingDigits"  # 训练数据集目录
+    test_data_dir = r"D:\二值手写数字数据集\testDigits"      # 测试数据集目录
+    
     # 加载训练数据集
-    train_data, train_labels = load_data("trainingDigits")  # 使用正确的函数名 load_data
+    train_data, train_labels = load_dataset(train_data_dir)  # 使用正确的函数名 load_dataset
     
     # 加载测试数据集
-    test_data, test_labels = load_data("testDigits")  # 加载测试集文件夹中的所有数据
+    test_data, test_labels = load_dataset(test_data_dir)  # 加载测试集文件夹中的所有数据
 
     # 测试集预测结果
     correct_count = 0
@@ -92,5 +95,3 @@ if __name__ == "__main__":
 
     accuracy = correct_count / total_count
     print(f"分类精度: {accuracy * 100:.2f}%")
-
-    
\ No newline at end of file