数据集筛选的目的

  我们获取到一个全新的成品数据集，未必所有类别都是我们所需要的。假如我们只需要几类数据来进行训练，就可以选择这种筛选方式

数据集筛选

• 以保留0，1，5，8类为例

import os
import shutil

# 文件夹路径
folder_path = "labels_test"


# 保留开头为0、1、5、8的行的函数
def filter_specific_lines(file_path):
    valid_starting_digits = {'0', '1', '5', '8'}
    temp_file = file_path + '.tmp'

    with open(file_path, 'r') as input_file, open(temp_file, 'w') as output_file:
        for line in input_file:
            if line.strip() and line[0] in valid_starting_digits:
                output_file.write(line)

    # 替换原文件
    shutil.move(temp_file, file_path)


# 处理文件夹中的所有文件
for filename in os.listdir(folder_path):
    file_path = os.path.join(folder_path, filename)
    if os.path.isfile(file_path):
        filter_specific_lines(file_path)

数据集映射重组

import os

# 定义类别映射关系，将0、1、5和8映射为0、1、2和3
class_mapping = {0: 0, 1: 1, 5: 2, 8: 3}

# 指定YOLO格式标注文件所在的文件夹路径
annotation_folder = '/path/to/annotations/folder'

# 循环处理每个标注文件
for filename in os.listdir(annotation_folder):
    annotation_file = os.path.join(annotation_folder, filename)
    if not annotation_file.endswith('.txt'):
        continue

    # 读取标注文件的内容
    with open(annotation_file, 'r') as f:
        lines = f.readlines()

    # 更新标注文件中的目标类别编号
    updated_lines = []
    for line in lines:
        parts = line.split()
        class_id = int(parts[0])
        if class_id in class_mapping:
            parts[0] = str(class_mapping[class_id])
            updated_lines.append(" ".join(parts))

    # 将更新后的标注信息保存回原始标注文件
    with open(annotation_file, 'w') as f:
        f.writelines(updated_lines)

根据重组后的txt文件筛选图片文件夹

import os

# 图片文件夹路径
image_folder_path = "1"
# txt文件夹路径
txt_folder_path = "labels_test_2"

# 获取所有txt文件的文件名，不包括扩展名
txt_files = set(os.path.splitext(filename)[0] for filename in os.listdir(txt_folder_path) if filename.endswith('.txt'))

# 遍历图片文件夹，删除没有对应txt文件的图片
for image_filename in os.listdir(image_folder_path):
    if image_filename.endswith('.jpg') or image_filename.endswith('.png'):
        image_name = os.path.splitext(image_filename)[0]
        if image_name not in txt_files:
            image_path = os.path.join(image_folder_path, image_filename)
            os.remove(image_path)
            print(f"Deleted {image_filename} because the corresponding txt file was not found.")

结尾

欢迎一起学习、讨论技术！
B站账号：Silver__Wolf_
Q：130856474