@Nikronic
Check this out:
class MasksTrainDataset(Dataset):
    """Masks Train dataset."""

    def __init__(self, csv_file, root_dir, transform=None):
        """
        Args:
            csv_file (string): Path to the csv file with annotations.
            root_dir (string): Directory with all the images.
            transform (callable, optional): Optional transform to be applied
                on a sample.
        """
        self.masks_frame = pd.read_csv(csv_file, skiprows=1500, nrows=10000)
        self.root_dir = root_dir
        self.transform = transform

    def __len__(self):
        return len(self.masks_frame)

    def __getitem__(self, idx):
        img_name = os.path.join(self.root_dir,
                                self.masks_frame.iloc[idx, 0])
        image = io.imread(img_name + '.png')
        image = cv2.cvtColor(image,cv2.COLOR_GRAY2RGB)
        ## use strip to get exact result
        masks = self.masks_frame.iloc[idx, 1].strip()
        if masks == '-1':
            mark = 0
        else:
            mark = 1
            
        #masks = np.array([masks])
        #masks = masks.astype('float').reshape(-1, 2)
        sample = {'image': image, 'masks': mark}

        if self.transform:
            sample = self.transform(sample)

        return sample
      
      
train_dataset = MasksTrainDataset(csv_file='pneumo_input/train/train-rle.csv',
                                    root_dir='pneumo_input/train/images/256/dicom/')
print(len(train_dataset))
for i in range(len(train_dataset)):
    sample = train_dataset[i]

    print(i, sample['image'].shape, sample['masks'])
    print(type(sample['masks']))

    if i == 4:
        plt.show()
        break

What do you think?