How make customised dataset for semantic segmentation?

I have two dataset folder of tif images, one is a folder called BMMCdata, and the other one is the mask of BMMCdata images called BMMCmasks(the name of images are corresponds). I am trying to make a customised dataset and also split the data randomly to train and test. Thank you in advance. at the moment I am getting an error

AttributeError: 'CustomDataset' object has no attribute 'len'

import torch
from torch.utils.data.dataset import Dataset  # For custom data-sets
from torchvision import transforms
from PIL import Image
import glob

folder_data = "D:\\Neda\\Pytorch\\U-net\\BMMCdata\\data"


class CustomDataset(Dataset):
    def __init__(self, root):   # initial logic happens like transform

        self.filename = folder_data
        self.root = root
        self.to_tensor = transforms.ToTensor()

        image_list = []
        for filename in glob.glob('folder_data/*.tif'):
            im = Image.open(filename)
            image_list.append(im)
            print(filename)

    def __getitem__(self, index):
        image = Image.open(self.image_list[index])
        return self.to_tensor(image)

    def __len__(self):  # return count of sample we have

        return self.len


custom_img = CustomDataset(folder_data)
# total images in set
print(custom_img.len)

train_len = int(0.6*custom_img.len)
test_len = custom_img.len - train_len
train_set, test_set = CustomDataset.random_split(custom_img, lengths=[train_len, test_len])
# check lens of subset
len(train_set), len(test_set)

train_set = CustomDataset(folder_data)
train_set = torch.utils.data.TensorDataset(train_set, train=True, batch_size=4)
train_loader = torch.utils.data.DataLoader(train_set, batch_size=4, shuffle=True, num_workers=1)
print(train_set)
print(train_loader)

test_set = torch.utils.data.DataLoader(Dataset, batch_size=4, sampler=test_set)
test_loader = torch.utils.data.DataLoader(Dataset, batch_size=4)

1 Like

There are some minor issues in your code:

  • The error is thrown, since you are returning self.len in __len__, without defining it. You should rather return something like len(self.image_list) (, which won’t currently work, but that’s the next point).
  • You are loading all images in __init__ without storing image_list. Your __getitem__ won’t be able to call self.image_list to open an image. However, even if you properly stored image_list using self.image_list = [], you are trying to open the image again. I would recommend to just store image paths in __init__ and load the images in __getitem__.
  • Currently you are not using any target. It seems just the data is loaded without the segmentation masks.
  • Dataset.random_split is not defined by default. You could split the image/mask paths and create two separate Datasets, one for training and another for validation.

Here is a small example I’ve written some time ago.
Let me know, if you can use it as a starter code and adapt it to your use case.

2 Likes

@ptrblck thanks a lot. I am reading your example. Will let you know.

1 Like

@ptrblck. I did manage to get this at the moment. I was expecting to get the count of data from len(). It’s returning me 35 for len, but I have 43 images and 43 masks.

Also, I add normalisation to the data as they are large values. To normalise them, I got the mean of an image (198) and standard deviation (64)? My question is:
should I get the standard deviation and mean of all dataset or one image?

Could you please let me know where I am doing wrong. thank you

I am getting an error train_set = CustomDataset(folder_mask) TypeError: __init__() missing 1 required positional argument: 'target_paths'

import torch
from torch.utils.data.dataset import Dataset  # For custom data-sets
import torchvision.transforms as transforms
from PIL import Image
import numpy as np
import torchvision
import matplotlib.pyplot as plt


folder_data = "D:\\Neda\\Pytorch\\U-net\\BMMCdata\\data"
folder_mask = "D:\\Neda\\Pytorch\\U-net\\BMMCmasks\\masks"


class CustomDataset(Dataset):
    def __init__(self, image_paths, target_paths, train=True):   # initial logic happens like transform

        self.image_paths = image_paths
        self.target_paths = target_paths
        self.transforms = transforms.Compose(
            [transforms.ToTensor(),
             transforms.Normalize((198, 198, 198), (64, 64, 64))])

    def __getitem__(self, index):

        image = Image.open(self.image_paths[index])
        mask = Image.open(self.target_paths[index])
        t_image, t_mask = self.transforms(image, mask)
        return t_image, t_mask

    def __len__(self):  # return count of sample we have

        return len(self.image_paths)


custom_img = CustomDataset(folder_data, folder_mask)
print(len(custom_img))

train_set = CustomDataset(folder_mask)
train_set = torch.utils.data.TensorDataset(train_set, train=True)
train_loader = torch.utils.data.DataLoader(train_set, batch_size=4, shuffle=True, num_workers=1)


test_set = CustomDataset(folder_data)
test_set = torch.utils.data.DataLoader(test_set, train=False)
test_loader = torch.utils.data.DataLoader(test_set, batch_size=4, shuffle=False, num_workers=1)


def imshow(img):
    img = img / 2
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))


dataiter = iter(train_loader)
images, mask = dataiter.next()

imshow(torchvision.utils.make_grid(images))
plt.show()

Currently you are just returning the length of the path, not the number of images.
image_paths should be a list of all paths to your images.
You can get all image paths using the file extension and a wildcard:

folder_data = glob.glob("D:\\Neda\\Pytorch\\U-net\\BMMCdata\\data\\*.jpg")
folder_mask = glob.glob("D:\\Neda\\Pytorch\\U-net\\BMMCdata\\masks\\*.jpg")

Change the extension to the right one for your files, e.g. .png, .tiff etc.

You could split these paths using a certain percentage:

len_data = len(folder_data)
train_size = 0.6
train_image_paths = folder_data[:int(len_data*train_size)]
test_image_paths = folder_data[int(len_data*train_size):]
train_mask_paths = folder_mask[:int(len_data*train_size)]
test_mask_paths = folder_mask[int(len_data*train_size):]

train_dataset = CustomDataset(train_image_paths, train_mask_paths)
test_dataset = CustomDataset(test_image_paths, test_mask_paths)

The mean and std values are usually calculated from all training images.
Also you should pass them as normalized values, since transforms.ToTensor() will normalize your data so that the values are in the range [0, 1].
A division of the mean and std by 255. should be sufficient.

In my example I’ve created a class function, which makes sure to apply the same random transformations on the data and target (e.g. RandomCrop). Since you are currently not using any random transformations, you could call your transformation on data and mask separately.
However, since you are dealing with a segmentation task, I doubt it’ll be a good idea to normalize your masks, since they should most likely contain the class indices.

def __getitem__(self, index):
    image = Image.open(self.image_paths[index])
    mask = Image.open(self.target_paths[index])
    x = self.transform(image)
    y = torch.from_numpy(np.array(mask)).long()
    return x, y

Let me know, if this would work for you!

2 Likes

@ptrblck Thank you very much for clarifying this. As I understood, the transforms.ToTensor() will normalize the value between [0, 1]. and I don’t need to normalize them? is that correct?

The new script looks like this and it’s returning the length of folder correctly as well. Could you please let me know if there is any mistake in this?

# get all the image and mask path and number of images
folder_data = glob.glob("D:\\Neda\\Pytorch\\U-net\\BMMCdata\\data\\*.tif")
folder_mask = glob.glob("D:\\Neda\\Pytorch\\U-net\\BMMCmasks\\masks\\*.tif")

# split these path using a certain percentage
len_data = len(folder_data)
print(len_data)
train_size = 0.6

train_image_paths = folder_data[:int(len_data*train_size)]
test_image_paths = folder_data[int(len_data*train_size):]

train_mask_paths = folder_mask[:int(len_data*train_size)]
test_mask_paths = folder_mask[int(len_data*train_size):]


class CustomDataset(Dataset):
    def __init__(self, image_paths, target_paths, train=True):   # initial logic happens like transform

        self.image_paths = image_paths
        self.target_paths = target_paths
        self.transforms = transforms.ToTensor()

    def __getitem__(self, index):

        image = Image.open(self.image_paths[index])
        mask = Image.open(self.target_paths[index])
        t_image = self.transforms(image)
        return t_image, mask

    def __len__(self):  # return count of sample we have

        return len(self.image_paths)


train_dataset = CustomDataset(train_image_paths, train_mask_paths, train=True)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=4, shuffle=True, num_workers=1)

test_dataset = CustomDataset(test_image_paths, test_mask_paths, train=False)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=4, shuffle=False, num_workers=1)

Yes, transforms.ToTensor will give you an image tensor with values in the range [0, 1].
This might be sufficient to train your model, however usually you would standardize your tensors to have zero-mean and a stddev of 1. You could calculate the mean and stddev of your train images yourself using this small example or alternatively the ImageNet mean and std work quite well for normal images.

The code looks good! You woun’t need the train argument anymore in your Dataset, but besides that I cannot find anything obviously wrong.

3 Likes

@ptrblck Thank you so much. You helped me to get a better understanding of customised database.

1 Like

Where are you getting the index from??

The index will load the images from the image path which stored in __init__

@ptrblck I saw a couple of example for schedule learning and save the best model, however most of them are implemented for transfer learning. I have hard time to generalise it to my own problem. I am following this tutorial for scheduling the learning rate and save the best model.

I split the previous dataset two three groups of train, validation and test and here is the code:

from custom_dataset import CustomDataset
import torch
import glob


# get all the image and mask path and number of images

folder_data = glob.glob("D:\\Neda\\Pytorch\\U-net\\my_data\\imagesResized\\*.png")
folder_mask = glob.glob("D:\\Neda\\Pytorch\\U-net\\my_data\\labelsResized\\*.png")

# split these path using a certain percentage
len_data = len(folder_data)
print("count of dataset: ", len_data)
# count of dataset:  992

split_1 = int(0.8 * len(folder_data))
split_2 = int(0.9 * len(folder_data))

folder_data.sort()

train_image_paths = folder_data[:split_1]
print("count of train images is: ", len(train_image_paths)) 
#count of train images is:  793

valid_image_paths = folder_data[split_1:split_2]
print("count of validation image is: ", len(valid_image_paths))
#count of validation image is:  99

test_image_paths = folder_data[split_2:]
print("count of test images is: ", len(test_image_paths)) 
#count of test images is:  100


#print(test_image_paths)

train_mask_paths = folder_mask[:split_1]

valid_mask_paths = folder_mask[split_1:split_2]

test_mask_paths = folder_mask[split_2:]


train_dataset = CustomDataset(train_image_paths, train_mask_paths)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=1, shuffle=True, num_workers=2)

valid_dataset = CustomDataset(valid_image_paths, valid_mask_paths)
valid_loader = torch.utils.data.DataLoader(valid_dataset, batch_size=1, shuffle=True, num_workers=2)

test_dataset = CustomDataset(test_image_paths, test_mask_paths)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=2)
  
DataLoader = {
        'train': train_loader,
        'valid': valid_loader,
        #'test': test_loader      
        }

when I follow the tutorial when it start for training at first epoch I am getting this error. could you please give me some suggestion how can I fix it? I assume is about DataLoader in above snippet.

for i, data in DataLoader[phase]:

ValueError: too many values to unpack (expected 2)

also, I am wondering how can I get the dataset_size[phase] in epoch_loss = running_loss / dataset_sizes[phase]
if I need to create a new thread please let me know.

Could you check the output of your Dataset?
Based on the example code in this thread it looks like two values should be returned, but I’m not sure how your current implementation works.
Just print the length of the returned values:

len(train_dataset[0])

You could get the size of the Dataset using this code snippet:

len(dataloaders[phase].dataset)

Also, as you can see, I’ve changed the name from DataLoader to dataloaders, since DataLoader is used by PyTorch and I don’t want to accidentally mask it somehow.

1 Like

this is current Dataset and it’s returning four values.

import torch
from torch.utils.data.dataset import Dataset  # For custom data-sets
import torchvision.transforms as transforms
from PIL import Image
import numpy 
import torchvision.transforms.functional
   
        
class CustomDataset(Dataset):
    def __init__(self, image_paths, target_paths):   # initial logic happens like transform

        self.image_paths = image_paths
        self.target_paths = target_paths
        self.transforms = transforms.ToTensor()
        self.mapping = {
            0: 0,
            255: 1              
        }
        
    def mask_to_class(self, mask):
        for k in self.mapping:
            mask[mask==k] = self.mapping[k]
        return mask
    
    def __getitem__(self, index):

        image = Image.open(self.image_paths[index])
        mask = Image.open(self.target_paths[index])
        t_image = image.convert('L')
        t_image = self.transforms(t_image)
        #mask = torch.from_numpy(np.array(mask))    #this is for BMCC dataset
        mask = torch.from_numpy(numpy.array(mask, dtype=numpy.uint8)) # this is for my dataset(lv)
        mask = self.mask_to_class(mask)
        mask = mask.long()
        return t_image, mask, self.image_paths[index], self.target_paths[index] 
    
    def __len__(self):  # return count of sample we have

        return len(self.image_paths)

ohhh, thank you for the dataloader.

In that case, you should assign all four returned value to a single one or to four separate ones:

for, data, target, data_path, target_path in dataloaders[phase]:
    ...

I think your code should work even if you use the same name, as it seems you didn’t import DataLoader directly. However, if it’s not too much of an effort, I would rename your dict.

1 Like

Thanks a lot. It works now but I think the training process is too slow in compare with before I did’t have validation and scheduling the learning rate. 4 epoch and step_size=2 completed in 7m. I will post a new thread for that as I have some questions about the scheduling the learning rate.

@ptrblck I am wondering how can I add a condition to CustomDataset for data augmentation only for few specific input images for training (image_207, image_387, image_502, image_508, image_509, image_520, image_597).

This is the CustomDataset snippet, basically, I added self.transformm to the previous code which posted above. I think I need to add a if condition in __getitem__ to apply self.transformm only on image. could you please point me in the right direction.

    
class CustomDataset(Dataset):
    def __init__(self, image_paths, target_paths, transform_images):   

        self.image_paths = image_paths
        self.target_paths = target_paths

        #self.aug = aug
        self.transformm = transforms.Compose([tf.rotate(10),
                                              tf.affine(0.2,0.2)])                                                                                   
        self.transform = transforms.ToTensor()
        
        self.transform_images = transform_images
            
        self.mapping = {
            0: 0,
            255: 1              
        }
        
    def mask_to_class(self, mask):
        for k in self.mapping:
            mask[mask==k] = self.mapping[k]
        return mask
    
    def __getitem__(self, index):

        image = Image.open(self.image_paths[index])
        mask = Image.open(self.target_paths[index])
        t_image = image.convert('L')
        t_image = self.transforms(t_image)
        
        if any([img in image for img in transform_images]):
            t_image = self.transformm(t_image)
                
        mask = torch.from_numpy(numpy.array(mask, dtype=numpy.uint8)) 
        mask = self.mask_to_class(mask)
        mask = mask.long()
        return t_image, mask, self.image_paths[index], self.target_paths[index] 
    
    def __len__(self):  # return count of sample we have

        return len(self.image_paths)

A simple approach would be to get the indices of these particular images and add a condition before applying the transformation.
Alternatively you could also check against self.image_paths:

def __init__(self, ..., transform_images):
    self.transform_images = transform_images
    ...

def __getitem__(self, index):
    ...
    if any([img in path for img in transform_images]):
        t_image = self.transformm(t_image)
        ...
1 Like

@ptrblck thank you. I still couldn’t manage do this. Could you please explain to me what is self.transform_images in __init__? where and how should I specify the indices of images (example image_207, image_208) that needs to be transform if I want to do the alternative solution?

self.transform_images would be a list containing all image names which should be transformed.
In __getitem__ the current image path will be checked against all image names in self.transform_images (the self. is missing in my code snippet).

I’m not sure, how you are creating the image paths, but once you get all the paths, you could use the .index() method on this list or alternatively use a condition to get the image indices.

1 Like

@ptrblck Thank you very much.