Pytorch数据集搭建与加载 pytorch-Datasets & DataLoaders +注释

import os import cv2 import torch import numpy as np from torch.utils.data import Dataset class MVSCustom(Dataset): # 获取数据路径和图片列表 def __init__(self,path): self.images_path = path self.get_image_lists() def get_image_lists(self): lists = os.listdir(self.images_path) #当前路径下所有文件名 lists = [r for r in lists if r[r.rfind("."):] == '.jpg'] #路径下除了图片可能还有其他文件，筛选出'.jpg'文件 self.rgb_paths = [os.path.join(self.images_path,r) for r in lists] # 图片的路径加上图片名 def read_image(self,filename): assert os.path.exists(filename), "file not found :{}".format(filename) # 如果，文件不存在，异常，结束 img = cv2.imread(filename) if img is None: # 有可能图像损坏，读出结果为空 print("!! read error =",filename) assert img is not None,"path is None" #img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB) np_img = np.array(img, dtype = np.float32) # 将numpy格式的图片，转为np.float32 return np_img def __len__(self): return len(self.rgb_paths) def __getitem__(self,index): rgb_path = self.rgb_paths[index] #路径中第index号图片 image = self.read_image(rgb_path) return image