74 lines
2.9 KiB
Python
74 lines
2.9 KiB
Python
# 输入数据集路径，划分完成训练集与测试集
|
||
# 配置文件里面进行路径拼接
|
||
|
||
from pathlib import Path
|
||
# import splitfolders
|
||
import os
|
||
|
||
# 读取，设置标签
|
||
def find_images_in_dir(directory, extensions=('.jpg', '.png', '.jpeg')):
|
||
for root, dirs, files in os.walk(directory):
|
||
for file in files:
|
||
if any(file.lower().endswith(ext) for ext in extensions):
|
||
yield os.path.join(root, file)
|
||
|
||
def path_subtract(full_path, base_path):
|
||
# 将两个路径都转换为绝对路径，以确保它们是可比较的
|
||
full_path = os.path.abspath(full_path)
|
||
base_path = os.path.abspath(base_path)
|
||
|
||
# 检查base_path是否是full_path的前缀
|
||
if full_path.startswith(base_path):
|
||
# 使用relpath函数获取相对于base_path的子路径
|
||
relative_path = os.path.relpath(full_path, base_path)
|
||
return relative_path
|
||
else:
|
||
# 如果base_path不是full_path的前缀，则不能相减（或者说结果无意义）
|
||
return None
|
||
|
||
def getlabel(root):
|
||
train_set = []
|
||
val_set = []
|
||
test_set = []
|
||
for image_path in find_images_in_dir(root):
|
||
# 处理每一个文件,训练集
|
||
subpath = path_subtract(image_path,root)
|
||
Parts = (Path(subpath)).parts[1]
|
||
|
||
if 'train' in subpath:
|
||
train_set.append([subpath,',',Parts.replace('_', ',') ])
|
||
if 'val' in subpath:
|
||
val_set.append([subpath,',',Parts.replace('_', ',') ])
|
||
if 'test' in subpath:
|
||
test_set.append([subpath,',',Parts.replace('_', ',') ])
|
||
|
||
if not os.path.exists(os.path.join(root,'label')):
|
||
# 如果文件夹不存在，则创建它
|
||
os.makedirs(os.path.join(root,'label'))
|
||
|
||
# 输出到txt
|
||
if len(train_set) > 0:
|
||
with open(os.path.join(root,'label','train.txt'), 'w', encoding='utf-8') as file:
|
||
# 使用循环将列表中的每一项写入文件，并在每项后面加上换行符
|
||
for item in train_set:
|
||
for sub in item:
|
||
file.write(sub)
|
||
file.write('\n')
|
||
if len(val_set) > 0:
|
||
with open(os.path.join(root,'label','val.txt'), 'w', encoding='utf-8') as file:
|
||
# 使用循环将列表中的每一项写入文件，并在每项后面加上换行符
|
||
for item in val_set:
|
||
for sub in item:
|
||
file.write(sub)
|
||
file.write('\n')
|
||
if len(test_set) > 0:
|
||
with open(os.path.join(root,'label','test.txt'), 'w', encoding='utf-8') as file:
|
||
# 使用循环将列表中的每一项写入文件，并在每项后面加上换行符
|
||
for item in test_set:
|
||
for sub in item:
|
||
file.write(sub)
|
||
file.write('\n')
|
||
|
||
if __name__ == '__main__':
|
||
# train:validation:test=8:1:1
|
||
getlabel('G:/dataset/test/split') |