Pytorch关于Dataset 的数据处理

from torch.utils.data import Dataset #Dataset的包
import os #路径需要这个
import cv2 # 需要读取图片，最好用opencv-python,当然也可以用PIL只是我不顺手


class MyData(Dataset): #我定义的这个类
    def __init__(self, root_dir, label_dir):
     #下面需要使用的变量，在__init__定义好，
        self.root_dir = root_dir # 根路径 data在电脑或者服务器大致的位置
        self.label_dir = label_dir # label的位置（这里假设Image的名字就是label的位置）
        self.path = os.path.join(self.root_dir, self.label_dir)# 将这个两个合在一起就能找到整体图片的大致路径
        self.img_path = os.listdir(self.path) #得到整体图片的路径（可取其中的一张一张的图像的名字）

    def __getitem__(self, idx): 
    # 改写__getitem__(self,item)函数，最后得到图像，标签
      #获取具体的一幅图像的名字
        img_name = self.img_path[idx]
        #获取一幅图像的详细地址
        img_item_path = os.path.join(self.root_dir, self.label_dir, img_name)
        #用opencv来读取图像
        img = cv2.imread(img_item_path)
        #获取标签（这里简单写了aligned与original）
        label = self.label_dir
        return img, label

    def __len__(self):
    #改写整体图像的大小
        return len(self.img_path)


root_dir = "D://data//basic"
img_dir = "Image"
aligned_label_dir = "aligned"
# aligned_label_dir = "Image//aligned"
aligned_label_dir = os.path.join(img_dir, aligned_label_dir)

original_label_dir = "original"
#original_label_dir = "Image//original"
original_label_dir = os.path.join(img_dir, original_label_dir)

#aligned_data = "D://data//basic//Image//aligned"
aligned_data = MyData(root_dir, aligned_label_dir)
#original_data = "D://data//basic//Image//original"
original_data = MyData(root_dir, original_label_dir)
data = aligned_data + original_data
# 15339
print(len(aligned_data))
# 15339
print(len(original_data))
# 30678
print(len(data))
img_1, label_1 = data[15338]
img_2, label_2 = data[15339]
print(label_1) # Image\aligned
print(label_2) # Image\original