LibTorch之DataSet数据集处理方法

$$LibTorch之DataSet数据集处理方法$$

C++读取文件夹下文件及文件夹路径

Pytorch

from torch.utils.data import Dataset
from PIL import Image
import os
# 返回所有图片路径和标签
def get_imgs_labels(data_dir):
    dict_label = {"ants": 0, "bees": 1}  # 如果改了分类目标，这里需要修改
    data_info = list()
    for root, dirs, _ in os.walk(data_dir):
        # 遍历类别
        for sub_dir in dirs:
            img_names = os.listdir(os.path.join(root, sub_dir))
            # img_names = list(filter(lambda x: x.endswith('.png'), img_names))   # 如果改了图片格式，这里需要修改
            # 遍历图片
            for i in range(len(img_names)):
                img_name = img_names[i]
                path_img = os.path.join(root, sub_dir, img_name)
                label = dict_label[sub_dir]
                data_info.append((path_img, int(label)))
    return data_info


class LoadDataset(Dataset):
    # 确定数据路径
    def __init__(self, data_dir=None, transform=None):
        self.imgs_labels = get_imgs_labels(data_dir)
        self.transform = transform

    # 获取并返回真实的数据和label
    def __getitem__(self, index):
        img_path,label = self.imgs_labels[index]
        img = Image.open(img_path)
        img.show()
        # print(type(img))
        # import numpy as np
        # img_array = np.array(img)
        # print(type(img_array))
        if self.transform is not None:
            img = self.transform(img)   # 在这里做transform，转为tensor等等
        return img, label

    # 确定索引的范围
    def __len__(self):
        return len(self.imgs_labels)



if __name__ == '__main__':
    import torchvision.transforms as transforms
    from torch.utils.data import DataLoader

    norm_mean = [0.33424968, 0.33424437, 0.33428448]
    norm_std = [0.24796878, 0.24796101, 0.24801227]

    train_transform = transforms.Compose([
        transforms.Resize((32, 32)),
        transforms.ToTensor(),  # 0-255 归一化到0-1 转Tensor
        transforms.Normalize(norm_mean, norm_std),
    ])

    train_dataset = LoadDataset(data_dir=r"D:\PycharmProjects\AI_Easy_Demo\MyData\split_data\train",transform=train_transform)

    train_loader = DataLoader(dataset=train_dataset, batch_size=8, shuffle=True)  # shuffle训练时打乱样本

    for idx,data_info in enumerate(train_loader):
        print(idx)
        inputs, labels = data_info
        print(data_info)

Libtorch

#include<opencv2/opencv.hpp>
#include <torch/torch.h>
#include <torch/script.h> 
#include <filesystem>
using namespace std;
namespace fs = std::filesystem;

vector<pair<string, int>> get_imgs_labels(const std::string& data_dir, map<string, int> dict_label)
{
	// 1.定义标签
	//map<string, int> dict_label;
	//dict_label.insert(pair<string, int>("ants", 0));
	//dict_label.insert(pair<string, int>("bees", 1));
	// 2.定义存储图像路径和标签的vector
	vector<pair<string, int>> data_info;
	// 3.读取图像和对应label放入data_info
	// 遍历字典,读取图像路径和对应label
	for (map<string, int>::iterator it = dict_label.begin(); it != dict_label.end(); it++)
	{
		// 遍历目录查找
		for (const auto& file_path : fs::directory_iterator(data_dir))
		{
			if (file_path.path().filename() == it->first) {
				// 遍历所有图像路径
				for (const auto& img_path : fs::directory_iterator(data_dir + "\\" + it->first))
				{
					//std::cout << img_path.path() << std::endl;
					data_info.push_back(pair<string, int>(img_path.path().string(), it->second));
				}
			}
			//std::cout << entry.path() << std::endl;
		}
		//printVector(data_info);
	}

	return data_info;
}

/// <summary>
/// 数据集处理模块类
/// </summary>
class MyDataset :public torch::data::Dataset<MyDataset> {
private:
	vector<pair<string, int>> data_info;
	torch::Tensor imgs, labels;
public:
	// 构造器：一般用于确定数据集和预处理形式
	MyDataset(const std::string& data_dir,std::map<string,int> dict_label);
	// get_item数据处理：对读取的数据进行预处理
	torch::data::Example<> get(size_t index) override;
	// 返回数据数量
	torch::optional<size_t> size() const override {
		return data_info.size();
	};
};

/// <summary>
/// 根据数据集路径和对应的标签列表，配对训练数据
/// </summary>
/// <param name="data_dir"></param>
/// <param name="dict_label"></param>
MyDataset::MyDataset(const std::string& data_dir, std::map<string, int> dict_label) {
	// 获取训练数据
	data_info = get_imgs_labels(data_dir, dict_label);

}

/// <summary>
/// 对数据进行预处理，并返回成对的实例Example{data,label}
/// </summary>
/// <param name="index"></param>
/// <returns></returns>
torch::data::Example<> MyDataset::get(size_t index)
{
	// 获取图像路径
	auto img_path = data_info[index].first;
    // 确定label
	auto label = data_info[index].second;
	// opencv根据图像路径读取图像
	auto image = cv::imread(img_path);
	// resize图像大小
	cv::resize(image, image, cv::Size(224, 224));
	// mat转tensor
	auto input_tensor = torch::from_blob(image.data, { image.rows, image.cols, 3 }, torch::kByte).permute({ 2, 0, 1 }).unsqueeze(0).to(torch::kFloat32) / 225.0;
	// int转tensor
	torch::Tensor label_tensor = torch::tensor(label);
	return {input_tensor,label_tensor };

}


int main()
{
	try
	{
		map<string, int> dict_label;
		dict_label.insert(pair<string, int>("ants", 0));
		dict_label.insert(pair<string, int>("bees", 1));
		// 设置dataset
		auto dataset_train = MyDataset("D:\\dataset\\hymenoptera_data\\train", dict_label).map(torch::data::transforms::Stack<>());
		// batchszie
		int batchSize = 1;
		// 设置dataloader
		auto dataLoader = torch::data::make_data_loader<torch::data::samplers::SequentialSampler>(std::move(dataset_train), batchSize);
		// 打印
		for (auto& batch : *dataLoader) {

			auto data = batch.data;
			auto target = batch.target;
			std::cout << data.sizes() << std::endl;
			//std::cout << data.max() << std::endl;
			//std::cout << data << std::endl;
			std::cout << target << std::endl;
			int ssss;
			cin >> ssss;
		}
	}
	catch (const std::exception& e)
	{
		// step5:打印报错
		cout << e.what() << endl;
	}
	
	return 0;
}