深度学习图像处理04：图像分类模型训练实战——动物分类

最新推荐文章于 2025-05-10 12:56:15 发布

White白小纯

最新推荐文章于 2025-05-10 12:56:15 发布

阅读量3.4k

点赞数 32

CC 4.0 BY-SA版权

分类专栏：深度学习深度学习图像处理 python 文章标签：深度学习人工智能 python

本文链接：https://blog.csdn.net/qq_35284513/article/details/137606961

这篇博文不涉及理论知识，主要通过一个完整的深度学习模型训练流程，直观地了解深度学习图像分类任务。有关理论的部分，之前几篇博文已经涉及基础部分，之后也会对一些理论进行补充。

本文将结合代码，主要介绍三部分内容：

1. 数据集划分

2. 模型训练

3. 模型评估

本文集成了深度学习图像处理从数据集划分到模型训练最后到模型评估的一个完整框架。旨在方便从事相关学术研究的伙伴，进行实验对比。

1. 数据集划分

1.1 数据集文件结构

数据集结构如上图所示，共分三级目录。根目录：dataset，一级子目录：animal-5，二级子目录：5种动物的类别。每个动物类别的文件夹中，保存对应类别的图片。

1.2 数据集划分

数据集划分的步骤：

（1）读取数据集中所有图片的路径。

（2）随机选取图片，将数据集按照4：1的比例，划分为训练集和测试集。

（3）标签映射，将数据集标签映射到json文件中保存。

1.3 代码

代码文件名称：split_data.py

完整代码：

import os
import random
import json
import argparse


def split_data(dataset_root, txt_save_root, ratio):
    assert 0 <= ratio <= 1, "ratio must be between 0 and 1"

    # 创建保存路径文件夹
    os.makedirs(txt_save_root, exist_ok=True)

    # 获取所有文件路径
    all_files = []
    for path, _, files in os.walk(dataset_root):
        for file in files:
            all_files.append(os.path.join(path, file))

    # 将文件路径保存到 dataset.txt
    with open(f'{txt_save_root}/dataset.txt', 'w') as f:
        for file_path in all_files:
            f.write(file_path + '\n')
    print(f'数据集路径保存成功到：{txt_save_root}/dataset.txt')

    # 随机分割为训练集和测试集
    random.shuffle(all_files)
    split_index = int(ratio * len(all_files))
    train_files = all_files[:split_index]
    test_files = all_files[split_index:]

    # 将训练集和测试集的路径分别保存
    with open(f'{txt_save_root}/train.txt', 'w') as f:
        for file_path in train_files:
            f.write(file_path + '\n')
    print(f'训练集路径保存成功到：{txt_save_root}/train.txt')

    with open(f'{txt_save_root}/test.txt', 'w') as f:
        for file_path in test_files:
            f.write(file_path + '\n')
    print(f'测试集路径保存成功到：{txt_save_root}/test.txt')


def get_paths_and_labels(filename):
    paths_and_labels = []
    with open(filename, 'r') as file:
        for line in file:
            path = line.strip()  # 移除行尾的换行符
            label = get_label_from_path(path)
            paths_and_labels.append((path, label))
    return paths_and_labels


def get_label_from_path(image_path):
    parts = image_path.split('\\')
    label = parts[-2] if len(parts) > 1 else None
    return label


def create_and_save_label_mapping(file_paths_and_labels, json_file):
    unique_labels = sorted(set(label for _, label in file_paths_and_labels))
    label_to_index = {label: index for index, label in enumerate(unique_labels)}

    with open(json_file, 'w') as file:
        json.dump(label_to_index, file)


if __name__ == '__main__':
    # 创建 ArgumentParser 对象
    parser = argparse.ArgumentParser(description="Process some paths.")
    # 添加参数
    parser.add_argument('--dataset_root', default='dataset/animal-5',
                        type=str, help='数据集根路径')
    parser.add_argument('--txt_save_root', default='file/animal-5',
                        type=str, help='txt文件保存路径')
    parser.add_argument('--ratio', default=0.8, type=float, help='训练集的比例')
    # 解析参数
    args = parser.parse_args()

    dataset_root = args.dataset_root
    txt_save_root = args.txt_save_root
    ratio = args.ratio

    # 划分数据集
    split_data(dataset_root, txt_save_root, ratio)
    # 获取数据集标签
    file_paths_and_labels = get_paths_and_labels(f'{txt_save_root}/train.txt')
    print(f'成功获取数据集标签')
    # 记录数据集标签映射
    create_and_save_label_mapping(file_paths_and_labels, f'{txt_save_root}/classes.json')
    print(f'已保存数据集标签映射到：{txt_save_root}/classes.json')

1.4 代码使用方法

parser.add_argument('--dataset_root', default='dataset/animal-5',
                    type=str, help='数据集根路径')
parser.add_argument('--txt_save_root', default='file/animal-5',
                    type=str, help='txt文件保存路径')
parser.add_argument('--ratio', default=0.8, type=float, help='训练集的比例')

运行代码时，只需要更改上述参数中的 default 部分。--dataset_root为数据集根路径。--txt_save_root为txt文件保存路径。--ratio为训练集占数据集的比例，大于0，小于1。

上述代码，运行成功后，结果如下：