跟着小土堆学习pytorch（一）——Dataset

一、前言

pytorch也是鸽了很久了，确定了下，还是用小土堆的教程。

kaggle获取数据集

二、dataset

dateset：数据集——提供一种方式去获取数据及其标签
如何获取数据及其标签以及总共多少个数据
dataloder：数据装载器——为后面的网络提供不同的数据形式
关于dataset也是重写类

三、代码展示

from torch.utils.data import Dataset
from PIL import Image
import os

class MyData(Dataset):
    def __init__(self, root_dir, label_dir):
        self.root_dir = root_dir
        self.label_dir = label_dir
        self.path = os.path.join(self.root_dir,self.label_dir)
        self.img_path = os.listdir(self.path)


    '''
    root_dir 根目录
    label_dir 路径
    os.path.join是为了把这两个路径加起来，用了这个函数后就不会出错了
    os.listdir这个是吧路径下的一系列文件名变成列表
    '''

    def __getitem__(self, idx):
        img_name = self.img_path[idx]
        img_item_path = os.path.join(self.root_dir,self.label_dir,img_name)
        img = Image.open(img_item_path)
        label = self.label_dir
        return img,label
    # 给getitem一个索引，就会返回img和label

    def __len__(self):
        return len(self.img_path)

root_dir = "dataset/train"
ants_label_dir = "ants"
bees_label_dir = "bees"
bees_dataset = MyData(root_dir,bees_label_dir)
ants_dataset = MyData(root_dir,ants_label_dir)

train_dataset = ants_dataset + bees_dataset
# 两个数据集的集合


# img_path = "D:\\pycharm_project\\try3\\dataset\\train\\ants\\0013035.jpg"
#
# image = cv2.imread(img_path)
#
# if image is not None:
#     # 如果成功读取，显示图片
#     cv2.imshow('Image', image)
#     cv2.waitKey(0)  # 等待按键
#     cv2.destroyAllWindows()  # 关闭所有窗口
# else:
#     print("Error: 图片读取失败。请检查图片路径。")