33 lines
949 B
Python
33 lines
949 B
Python
from torch.utils.data import Dataset
|
||
import cv2
|
||
from PIL import Image
|
||
import os
|
||
|
||
|
||
class MyDataset(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)
|
||
|
||
# __getitem__ 是一个魔法方法，当使用dataset[i]获取第i个样本时，就会调用该方法。
|
||
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
|
||
|
||
def __len__(self):
|
||
return len(self.img_path)
|
||
|
||
|
||
root_dir = 'dataset/train'
|
||
ants_label_dir = 'ants'
|
||
bees_label_dir = 'bees'
|
||
ants_dataset = MyDataset(root_dir, ants_label_dir)
|
||
bees_dataset = MyDataset(root_dir, bees_label_dir)
|
||
print(ants_dataset[0]) |