Pytorch creates its own dataset and loads the dataset

Create a class and inherit torch.utils.data.dataset.Datase class

class MyDataset(Dataset):
    ''' data_path:  Dataset path  img_size:  Picture size  train_lines:  Image name array  '''
    def __init__(self,data_path,img_size,train_lines):
        super(MyDataset, self).__init__()

        self.data_path = data_path
        self.img_size = img_size
        self.train_lines = train_lines
        self.length = len(train_lines)

establish __getitem__ Method

class MyDataset(Dataset):
    ''' data_path:  Dataset path  img_size:  Picture size  train_lines:  Image name array  '''
    def __init__(self,data_path,img_size,train_lines):
        super(MyDataset, self).__init__()

        self.data_path = data_path
        self.img_size = img_size
        self.train_lines = train_lines
        
    def __getitem__(self, index):
        annotation_line = self.train_lines[index]
        name = annotation_line.split()[0]       #  Get the picture name 

        image = Image.open(os.path.join(os.path.join(self.data_path,"dem"),name+".tif"))
        label = Image.open(os.path.join(os.path.join(self.data_path, "label"), name + ".png"))

        image = np.array(image)
        label = np.array(label)

        image = cv2.resize(image,(self.img_size,self.img_size))
        label = cv2.resize(label,(self.img_size,self.img_size))

        # image = image[np.newaxis,:] 

        print("images size: {}, label size: {}".format(image.shape,label.shape))

        return image,label

Load data set

If you don't know how to write the names of all pictures in the folder TXT We can refer to ：python Read all the pictures in the folder and write the picture names line by line txt in ：https://blog.csdn.net/weixin_43598687/article/details/125666776?spm=1001.2014.3001.5501

dataset_path = r"E:/workspace/PyCharmProject/dem_feature/dem/512"

#  Open the of the dataset txt,  Read the picture name line by line 
with open(os.path.join(dataset_path, "dem/train.txt"), "r") as f:
    train_lines = f.readlines()

with open(os.path.join(dataset_path, "dem/val.txt"), "r") as f:
    val_lines = f.readlines()

train_dataset = MyDataset(dataset_path, img_size=512,train_lines=train_lines)

train_dataloader = DataLoader(train_dataset,batch_size=8,shuffle=False)

for iteration,data in enumerate(train_dataloader):
    imgs,labels = data
    
    print(imgs,labels)