制作数据集简单教程

制作数据集简单教程

简介

​ 深度学习神经网络模型发展迅速，下游应用众多，预训练模型已经能够精准识别多种目标，但是当我们需要按需精细化检测自己的目标时还是需要自己收集制作数据集。

1.开源数据集

​ 去各种数据集网站下载开源的数据集，数据集标注文件可能需要转换格式。

Kaggle: https://www.kaggle.com/

Dataset Search: https://datasetsearch.research.google.com/

五号雷达: https://www.5radar.com/

Finddata: https://www.findata.cn/

ModelScope: https://modelscope.cn/datasets

2.自制数据集

2.1 爬取图像

​ 通过各种搜索引擎搜索图像，收集自己需要的数据，但是效率较低，因此一般使用爬虫程序(需要注意版权问题)下载目标图像后对收集的数据进行清洗。
​ 例如想快速爬取百度图片，需要从百度图片搜索获取中间尺寸图片的URL列表，函数的返回值是一个包含图片URL的列表。使用正则表达式提取middleURL字段来获取图片地址。通过分页参数page实现分页搜索，每次请求返回图片链接。示例函数(传入搜索关键字和搜索结果页数参数，返回图片链接列表)如下：

def pic_req(keyword, page):
    #请求头，模拟真实用户浏览器
    headers = {
    "Accept": "text/plain, */*; q=0.01",
    "Accept-Language": "zh-CN,zh;q=0.9,en;q=0.8,ja;q=0.7",
    "Cache-Control": "no-cache",
    "Connection": "keep-alive",
    "Pragma": "no-cache",
    "Referer": "https://image.baidu.com/search/index...",
    "Sec-Fetch-Dest": "empty",
    "Sec-Fetch-Mode": "cors",
    "Sec-Fetch-Site": "same-origin",
    "User-Agent": "Mozilla/5.0...",
    "X-Requested-With": "XMLHttpRequest",
    "sec-ch-ua": "\"Not/A)Brand\";v=\"8\", \"Chromium\"...",
    "sec-ch-ua-mobile": "?0",
    "sec-ch-ua-platform": "\"Windows\""
}

​ cookies:有助于通过百度验证和获取数据

cookies = {
    "RT": "\"z=1&dm=baidu.com...",
    "BIDUPSID": "C59FC4562EF36BE1EA24CB3C40296780",
    "PSTM": "1718722340",
    ...
}

​ URL(百度图片搜索的API端点)和请求参数parms(关键字和分页参数)

url = "https://image.baidu.com/search/acjson" 
params = {
    "tn": "resultjson_com",
    "logid": "11535577539341025444",
    "ipn": "rj",
    "ct": "201326592",
    "is": "",
    "fp": "result",
    "word": keyword,
    "queryWord": keyword,
    ...
    "pn": 30*page,  # 设置请求的页数
    "rn": "30",  # 每页返回30条结果
}

​ 发送请求并响应处理

try:
    # 使用requests.get发送GET请求，获取响应内容并将其转换为纯文本
    response = requests.get(url, headers=headers, cookies=cookies, params=params).text
    response = response.replace("\\", "")
    # 通过正则表达式pattern匹配响应中的图片URL，返回符合格式的图片链接列表
    pattern = r'"middleURL":"(https?://[^"]+)"'
    matches = re.findall(pattern, response)
    return matches
# 异常处理：捕获请求或JSON解析失败的异常并输出错误信息
except requests.exceptions.RequestException as e:
    print(f"请求失败: {e}")
    return []
except ValueError as e:
    print(f"JSON解码失败: {e}")
    return []

​ 存储图片，保存到 images 目录下并以URL中的文件名命名：

def pic_store(url):
    try:
        img_data = requests.get(url).content
        with open(f"images/{url.split('/')[-1]}", 'wb') as handler:
            handler.write(img_data)
        print(f"图片已保存: {url}")
    except Exception as e:
        print(f"图片保存失败: {e}")

​ 只需要在主函数中给出搜素关键词的变量，建立一个循环调用爬取函数，再将图片存储为需要的格式即可

if __name__ == "__main__":
    keyword = "抽屉柜"
    for i in range(1, 10):  # 爬取页数
        urls = pic_req(keyword, i)
        for url in urls:
            pic_store(url)

2.2 自主拍摄

​ 应用在特定场景的检测需要收集目标场景下的数据，以增强模型的实用性和可靠性，在实际工作环境下拍摄图片或者视频，对其进行标注，以获取高质量数据集。

​ 视频截取图片python程序(传入视频路径、输出文件夹名字以及提取帧率)示例如下：

import cv2
import os

def extract_frames(video_path, output_folder, frame_rate=1):
    # 检查输出文件夹是否存在，如果不存在则创建
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    # 打开视频文件
    cap = cv2.VideoCapture(video_path)
    # 获取视频帧率
    video_fps = cap.get(cv2.CAP_PROP_FPS)
    # 计算帧间隔，即每隔多少帧保存一帧
    interval = int(video_fps / frame_rate)
    frame_count = 0 # 记录当前视频序号
    saved_frame_count = 0 # 记录保存的帧数
    
    # 读取视频帧
    while cap.isOpened():
        ret, frame = cap.read() # 读取一帧
        if not ret: # 读取失败则跳出循环
            break
        if frame_count % interval == 0: # 间隔保存
            # 生成保存的帧文件名，按顺序编号，格式为 "frame_0000.jpg"
            frame_filename = os.path.join(output_folder, f'frame_{saved_frame_count:04d}.jpg')
            # 将帧保存为图片
            cv2.imwrite(frame_filename, frame)
            saved_frame_count += 1
        frame_count += 1
    cap.release()
    # 输出保存帧的总数和路径
    print(f"Extracted {saved_frame_count} frames and saved to {output_folder}")

# 主函数
video_path = '/home/yds/mydatasets/scripts/captured_videos/video.avi'
output_folder = 'pic'
extract_frames(video_path, output_folder, frame_rate=1)  # 每秒截取一帧

2.3 数据预处理

统一大小

​ 收集的图片有可能尺寸大小不同，将其规格统一：

# 修改成图片所需的宽和高，循环调用即可
width = 640 
height = 480
img = cv2.resize(img, (width, height), interpolation=cv2.INTER_AREA)

数据增强

​ 现有数据不足时可以使用数据增强生成更多数据，常见方式有几何、颜色和亮度变换等。

# 水平变换
def pingyi(root_path, img_name):
    img = Image.open(os.path.join(root_path, img_name))  # 打开图像文件
    img = cv2.cvtColor(numpy.asarray(img), cv2.COLOR_RGB2BGR)
    cols, rows = img.shape[0], img.shape[1]
    # 创建平移矩阵，水平平移50像素，垂直平移30像素
    M = np.float32([[1, 0, 50], [0, 1, 30]])
    # 应用平移变换
    dst = cv2.warpAffine(img, M, (cols, rows), borderValue=(0, 255, 0))
    pingyi_img = Image.fromarray(cv2.cvtColor(dst, cv2.COLOR_BGR2RGB))
    return pingyi_img

# 旋转变换
def rotation(root_path, img_name):
    img = Image.open(os.path.join(root_path, img_name))
    # 随机选择旋转角度，范围为-2到2的90度倍数，若随机角度为0，则旋转-90度
    random_angle = np.random.randint(-2, 2) * 90
    if random_angle == 0:
        rotation_img = img.rotate(-90)
    else:
        # 使用expand=True，确保旋转后图像的边界自适应调整
        rotation_img = img.rotate(random_angle, expand=True)
    return rotation_img

# 水平翻转
def flip(root_path, img_name):
    img = Image.open(os.path.join(root_path, img_name))
    filp_img = img.transpose(Image.FLIP_LEFT_RIGHT) # 进行水平翻转
    return filp_img

# 亮度增强
def brightnessEnhancement(root_path, img_name):
    image = Image.open(os.path.join(root_path, img_name))
    # 创建亮度增强对象
    enh_bri = ImageEnhance.Brightness(image)
    # 设置亮度增强的比例，这里固定为1.5倍
    brightness = 1.5
    # 应用亮度增强
    image_brightened = enh_bri.enhance(brightness)
    return image_brightened

# 对比度增强
def contrastEnhancement(root_path, img_name):
    image = Image.open(os.path.join(root_path, img_name))
    # 创建对比度增强对象
    enh_con = ImageEnhance.Contrast(image)
    # 设置对比度增强的比例
    contrast = 1.5
    image_contrasted = enh_con.enhance(contrast)
    return image_contrasted

​ 对图片数据随机使用数据增强方法：

# 随机组合增强方式的函数
def apply_random_combination(root_path, img_name, save_path):
    # 定义增强方法的列表
    methods = [pingyi, rotation, flip, brightnessEnhancement, contrastEnhancement]
    # 随机选择至少一种增强方式（可以选择多种）
    selected_methods = random.sample(methods, random.randint(1, len(methods)))
    # 打开图像
    img = Image.open(os.path.join(root_path, img_name))
    # 随机选择的增强方法，对文件夹中图片逐一应用
    for method in selected_methods:
        img = method(root_path, img_name)
    # 保存增强后的图像
    save_name = f"enhanced_combined_{img_name}" 
    img.save(os.path.join(save_path, save_name))

2.4 标注数据

​ 标注自己想要识别定位的目标物体，让计算机通过图像学会识别和分类各种目标。监督学习必须要标注，通过学习现有的数据理解目标的特征和不同类别物体的差异。可以使用LabelImg、Labelme、Make Sense等标注工具进行标注，下面以LabelImg为例。
​ 安装标注工具LabelImg：

pip install labelimg

​ 安装完成后输入：

labelImg

​ 出现标注界面：

​ 打开 Open Dir 选择存放图像数据的路径，打开Change Save Dir 选择对应的标签路径(默认存放在图像路径下)，数据集格式改为 YOLO 点击左上角的view勾选自动保存和展示标签，按 W 创建一个新的标注框，通过鼠标拖动、调整大小和位置等方式，精确地绘制，完成之后按 D 键进入下一张图片。
​ 使用标注工具标记图像后，标签导出为.txt文件，每张图像对应一个同名的.txt 文件(如果图像中没有对象，则不需要)。.txt 文件规格如下：
①每个对象一行；
②每一行都是class x_center y_center width height的数据格式；
③检测框坐标x y w h(目标的坐标值 x、y，检测框的宽 w 和高 h)的数据范围在0-1间。如果检测框是以像素为单位，将x_center 和 width 除以图像宽度，将y_center 和 height 除以图像高度；
④类编号从索引0开始。

​ 示例图片的标签文件如下：

​ 该标签文件表示图像中标注有2个人(所属类别定义为 0)和一根领带(所属类别定义为 27)。标注完成后将所有数据(图像和标签同名，一一对应)按 8：2 的比例划分为训练集和验证集(测试集可为空)，放入前面创建的文件夹中，构建数据集完毕。

2.5 划分数据集

需要将标注好的数据集按比例划分为训练集和验证集，训练集用于训练神经网络模型的参数和权重，通过多次迭代优化提高性能；验证集用于调整模型超参数并初步评估模型，防止过拟合。划分数据集示例代码如下：

import os
import shutil
import random
import os.path as osp

# 将图像和标签文件夹中的数据集按比例划分为训练集和测试集,并保存到指定文件夹中
def split_datasets(image_folder, label_folder, output_folder, train_ratio=0.7, test_ratio=0.3, random_seed=None):
    """
    参数:
    - image_folder: 输入图像文件夹路径
    - label_folder: 输入标签文件夹路径
    - output_folder: 输出文件夹的根路径,将在其中创建两个子文件夹:train和test,每个子文件夹中有images和labels
    - train_ratio: 训练集比例
    - test_ratio: 测试集比例
    - random_seed: 随机数种子，可选
    """
    
    # 获取图像和标签文件列表并排序
    images = sorted(os.listdir(image_folder))
    labels = sorted(f for f in os.listdir(label_folder) if f != 'classes.txt')

    # 确保图像和标签数量一致
    assert len(images) == len(labels), "图像和标签文件数量不一致"

    # 确保图像和标签文件名对应
    for img, lbl in zip(images, labels):
        assert osp.splitext(img)[0] == osp.splitext(lbl)[0], "文件名不对应: {} 和 {}".format(img, lbl)

    # 组合并打乱图像和标签文件
    combined = list(zip(images, labels))
    random.seed(random_seed)
    random.shuffle(combined)

    num_files = len(combined)
    num_train = int(num_files * train_ratio)
    num_test = num_files - num_train

    train_files = combined[:num_train]
    test_files = combined[num_train:num_train + num_test]

    # 创建输出文件夹
    train_image_folder = osp.join(output_folder, 'Train', 'images')
    train_label_folder = osp.join(output_folder, 'Train', 'labels')
    test_image_folder = osp.join(output_folder, 'Test', 'images')
    test_label_folder = osp.join(output_folder, 'Test', 'labels')

    os.makedirs(train_image_folder, exist_ok=True)
    os.makedirs(train_label_folder, exist_ok=True)
    os.makedirs(test_image_folder, exist_ok=True)
    os.makedirs(test_label_folder, exist_ok=True)

    # 将文件复制到相应的文件夹
    for img, lbl in train_files:
        shutil.copy(osp.join(image_folder, img), osp.join(train_image_folder, img))
        shutil.copy(osp.join(label_folder, lbl), osp.join(train_label_folder, lbl))
    for img, lbl in test_files:
        shutil.copy(osp.join(image_folder, img), osp.join(test_image_folder, img))
        shutil.copy(osp.join(label_folder, lbl), osp.join(test_label_folder, lbl))

    # 复制 classes.txt 文件到每个子文件夹
    classes_txt = osp.join(label_folder, 'classes.txt')
    if osp.exists(classes_txt):
        for subfolder in ['Train', 'Test']:
            os.makedirs(osp.join(output_folder, subfolder, 'labels'), exist_ok=True)
            shutil.copy(classes_txt, osp.join(output_folder, subfolder, 'labels', 'classes.txt'))

# 使用
if __name__ == "__main__":
    image_folder = "/home/yds/mydatasets/make_data/cabinet_add/images"
    label_folder = "/home/yds/mydatasets/make_data/cabinet_add/labels"
    output_folder = "/home/yds/mydatasets/make_data/cabinet_add"
    
    split_datasets(image_folder, label_folder, output_folder, train_ratio=0.8, test_ratio=0.2, random_seed=42)

​ 划分好的数据集文件夹格式(COCO128)如下：

​ 划分完毕后即可得到用于训练模型的数据集。