把这段代码改写成可随时调用的函数,功能不变
时间: 2024-10-19 16:03:59 浏览: 23
JS自调用匿名函数具体实现
当然可以,你可以将这段代码封装成一个名为`create_dataset_csv`的函数,以便在需要的时候随时调用。下面是一种可能的实现:
```python
import os
import pandas as pd
def create_dataset_csv(dataset_dir, classes, output_file='dogvscat_dataset.csv'):
"""
创建包含文件信息的CSV文件
参数:
- dataset_dir (str): 数据集的根目录
- classes (list): 类别名称列表,对应目录里的子目录
- output_file (str): 输出的CSV文件名,默认为'dogvscat_dataset.csv'
返回:
None, 直接将数据写入csv文件
"""
col = ['filename', 'filepath', 'label']
info_array = []
for kindname in classes:
classpath = os.path.join(dataset_dir, kindname)
for filename in os.listdir(classpath):
filepath = os.path.join(classpath, filename)
label = classes.index(kindname)
info_array.append([filename, filepath, label])
info_array = np.array(info_array)
df = pd.DataFrame(info_array, columns=col)
df.to_csv(output_file, encoding='utf-8')
# 调用函数
create_dataset_csv('your_dataset_directory', ['dog', 'cat'])
```
现在你可以通过传入适当的参数来调用`create_dataset_csv`函数,例如指定数据集目录和类别列表,以及想要生成的CSV文件名。如果不需要返回值,直接调用即可。
阅读全文