维基百科中爬取的数据集:
代码如下:
import os
import requests
from bs4 import BeautifulSoup
# 设置要爬取的维基百科页面
wiki_page = "https://en.wikipedia.org/wiki/ImageNet"
# 创建用于保存图片的目录
dataset_dir = "wiki_image_dataset"
os.makedirs(dataset_dir, exist_ok=True)
# 发送 HTTP 请求获取页面内容
response = requests.get(wiki_page)
html_content = response.content
# 使用 BeautifulSoup 解析 HTML 内容
soup = BeautifulSoup(html_content, "html.parser")
# 查找页面中的所有图片链接
img_tags = soup.find_all("img")
# 遍历图片链接,下载并保存图片
for i, img_tag in enumerate(img_tags):
img_src = img_tag.get("src")
if img_src.startswith("//"):
img_src = "https:" + img_src
# 构建图片文件名
img_filename = os.path.join(dataset_dir, f"image_{i}.jpg")
# 下载并保存图片
img_response = requests.get(img_src)
with open(img_filename, "wb") as f:
f.write(img_response.content)
print(f"Downloaded and saved image: {img_filename}")
print(f"Total {len(img_tags)} images have been saved to the '{dataset_dir}' directory.")
将维基百科中的数据集划分为训练集测试集验证集合
代码如下:
import random
from shutil import copy2
def data_set_split(src_data_folder, target_data_folder, train_scale=0.8, val_scale=0.0, test_scale=0.2):
'''
读取源数据文件夹,生成划分好的文件夹,分为trian、val、test三个文件夹进行
:param src_data_folder: 源文件夹
:param target_data_folder: 目标文件夹
:param train_scale: 训练集比例
:param val_scale: 验证集比例
:param test_scale: 测试集比例
:return:
'''
print("开始数据集划分")
class_names = os.listdir(src_data_folder)
# 在目标目录下创建文件夹
split_names = ['train', 'val', 'test']
for split_name in split_names:
split_path = os.path.join(target_data_folder, split_name)
if os.path.isdir(split_path):
pass
else:
os.mkdir(split_path)
# 然后在split_path的目录下创建类别文件夹
for class_name in class_names:
class_split_path = os.path.join(split_path, class_name)
if os.path.isdir(class_split_path):
pass
else:
os.mkdir(class_split_path)
# 按照比例划分数据集,并进行数据图片的复制
# 首先进行分类遍历
for class_name in class_names:
current_class_data_path = os.path.join(src_data_folder, class_name)
current_all_data = os.listdir(current_class_data_path)
current_data_length = len(current_all_data)
current_data_index_list = list(range(current_data_length))
random.shuffle(current_data_index_list)
train_folder = os.path.join(os.path.join(target_data_folder, 'train'), class_name)
val_folder = os.path.join(os.path.join(target_data_folder, 'val'), class_name)
test_folder = os.path.join(os.path.join(target_data_folder, 'test'), class_name)
train_stop_flag = current_data_length * train_scale
val_stop_flag = current_data_length * (train_scale + val_scale)
current_idx = 0
train_num = 0
val_num = 0
test_num = 0
for i in current_data_index_list:
src_img_path = os.path.join(current_class_data_path, current_all_data[i])
if current_idx <= train_stop_flag:
copy2(src_img_path, train_folder)
# print("{}复制到了{}".format(src_img_path, train_folder))
train_num = train_num + 1
elif (current_idx > train_stop_flag) and (current_idx <= val_stop_flag):
copy2(src_img_path, val_folder)
# print("{}复制到了{}".format(src_img_path, val_folder))
val_num = val_num + 1
else:
copy2(src_img_path, test_folder)
# print("{}复制到了{}".format(src_img_path, test_folder))
test_num = test_num + 1
current_idx = current_idx + 1
print("*********************************{}*************************************".format(class_name))
print(
"{}类按照{}:{}:{}的比例划分完成,一共{}张图片".format(class_name, train_scale, val_scale, test_scale,
current_data_length))
print("训练集{}:{}张".format(train_folder, train_num))
print("验证集{}:{}张".format(val_folder, val_num))
print("测试集{}:{}张".format(test_folder, test_num))
if __name__ == '__main__':
src_data_folder = r"D:\Chinese Medicine" # todo 原始数据集目录
target_data_folder = r"D:\1" # todo 数据集分割之后存放的目录
data_set_split(src_data_folder, target_data_folder)