71 lines
2.5 KiB
Python
71 lines
2.5 KiB
Python
from torch.utils.data import Dataset, DataLoader
|
||
import numpy as np
|
||
from PIL import Image
|
||
import os
|
||
from torchvision import transforms
|
||
from torch.utils.tensorboard import SummaryWriter
|
||
from torchvision.utils import make_grid
|
||
|
||
writer = SummaryWriter("logs")
|
||
|
||
class MyData(Dataset):
|
||
|
||
def __init__(self, root_dir, image_dir, label_dir, transform):
|
||
self.root_dir = root_dir
|
||
self.image_dir = image_dir
|
||
self.label_dir = label_dir
|
||
self.label_path = os.path.join(self.root_dir, self.label_dir)
|
||
self.image_path = os.path.join(self.root_dir, self.image_dir)
|
||
self.image_list = os.listdir(self.image_path)
|
||
self.label_list = os.listdir(self.label_path)
|
||
self.transform = transform
|
||
# 因为label 和 Image文件名相同，进行一样的排序，可以保证取出的数据和label是一一对应的
|
||
self.image_list.sort()
|
||
self.label_list.sort()
|
||
|
||
def __getitem__(self, idx):
|
||
img_name = self.image_list[idx]
|
||
label_name = self.label_list[idx]
|
||
img_item_path = os.path.join(self.root_dir, self.image_dir, img_name)
|
||
label_item_path = os.path.join(self.root_dir, self.label_dir, label_name)
|
||
img = Image.open(img_item_path)
|
||
|
||
with open(label_item_path, 'r') as f:
|
||
label = f.readline()
|
||
|
||
# img = np.array(img)
|
||
img = self.transform(img)
|
||
sample = {'img': img, 'label': label}
|
||
return sample
|
||
|
||
def __len__(self):
|
||
assert len(self.image_list) == len(self.label_list)
|
||
return len(self.image_list)
|
||
|
||
if __name__ == '__main__':
|
||
transform = transforms.Compose([transforms.Resize((256, 256)), transforms.ToTensor()])
|
||
root_dir = "dataset/train"
|
||
image_ants = "ants_image"
|
||
label_ants = "ants_label"
|
||
ants_dataset = MyData(root_dir, image_ants, label_ants, transform)
|
||
image_bees = "bees_image"
|
||
label_bees = "bees_label"
|
||
bees_dataset = MyData(root_dir, image_bees, label_bees, transform)
|
||
train_dataset = ants_dataset + bees_dataset
|
||
|
||
# transforms = transforms.Compose([transforms.Resize(256, 256)])
|
||
dataloader = DataLoader(train_dataset, batch_size=1, num_workers=2)
|
||
|
||
writer.add_image('error', train_dataset[119]['img'])
|
||
writer.close()
|
||
for i, j in enumerate(dataloader):
|
||
# imgs, labels = j
|
||
print(type(j))
|
||
print(i, j['img'].shape)
|
||
# writer.add_image("train_data_b2", make_grid(j['img']), i)
|
||
|
||
writer.close()
|
||
|
||
|
||
|