Pytorch 图片的读取

• 图片的读取

from PIL import Image
img = Image.open(path)
# 这样读取到的图片格式是PIL格式

或者

import cv2
img = cv2.imread(path)
# 这样读取的图片是np.array()格式

• 自定义函数读取图片并储存到列表

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

class MyData(Dataset):
    def __init__(self, root_dir, label_dir):
        # 图片的label为文件夹名称
        # self. 使变量成为这个类中的全局变量
        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)
        # os.listdir 返回指定路径下的文件夹和文件列表

    def __getitem__(self, idx):
        # 第i张图片的名称
        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
    
    del __len__(self):
        # 有几张图片
        return len(self.img_path)

要调用它只需

root_dir = "data/train"
label_dir = "ants"
ants_dataset = MyData(root_dir, label_dir)
img, label = ants_dataset[0]