Create custom dataset from multiple folders

Crispolo: <p>I have some images organized in folders as shown in the following picture:</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/9/3/93a52b9e99dba803a142c568a978e276c727d0e7.png" alt="image" data-base62-sha1="l483YkwOByhteEBC8WogHNIwsp9" width="170" height="314"></p>
<p>In order to create a PyTorch DataLoader I defined a custom Dataset in this way</p>
<pre><code class="lang-auto">class CustomDataset(Dataset):
  def __init__(self, root, dirs=None, transforms=None):
    self.root_dir = root
    self.sel_dirs = dirs
    self.transforms = transforms

  def __getFileNr(self, path):
    count = 0

    for f in os.listdir(path):
      absPathIn = os.path.join(path, f)

      if os.path.isdir(absPathIn):
        count += self.__getFileNr(absPathIn)  
      elif os.path.isfile(absPathIn):
        count += 1

    return count

  def __len__(self):
    if self.sel_dirs == None:
      return

    total_count = 0
    for dir in self.sel_dirs:
      path = os.path.join(self.root_dir, dir)

      if os.path.isfile(path):
        total_count += 1
      elif os.path.isdir(path):
        count = 0
        count += self.__getFileNr(path)
        total_count += count

    return total_count

  def __checkDir(self, dirPath):
    images = []
    labels = []

    for f in os.listdir(dirPath):
      absDirPath = os.path.join(dirPath, f)

      if os.path.isdir(absDirPath):
        tmpImgs, tmpLabels = self.__checkDir(absDirPath)
        images.extend(tmpImgs)
        labels.extend(tmpLabels)

      elif os.path.isfile(absDirPath):
        images.append(absDirPath)
        fileName = os.path.splitext(os.path.split(absDirPath)[1])[0]
        labels.append(fileName.split('_')[-1])
        
    return (images, labels)

  def __getitem__(self):
    images = []
    labels = []
        
    for dir in self.sel_dirs:
      dirPath = os.path.join(self.root_dir, dir)
 
      if os.path.isdir(dirPath):
        tmpImgs, tmpLabels = self.__checkDir(dirPath)
        images.extend(tmpImgs)
        labels.extend(tmpLabels)
      else:
        return

    if self.transforms and len(images) &gt; 0:
      images = self.transforms(images)

    return images, labels
</code></pre>
<p>Everything seems to work, I’m new to PyTorch and I’m not sure this is a good solution, however when I try to create the DataLoader, I obtain the following error message:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/a/8a3d4c1889c09b3c6680adcf41ca8ff02d9b4190.png" data-download-href="https://discuss.pytorch.org/uploads/default/8a3d4c1889c09b3c6680adcf41ca8ff02d9b4190" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/8/a/8a3d4c1889c09b3c6680adcf41ca8ff02d9b4190.png" alt="image" data-base62-sha1="jIVdCqgCth8atA7AbRs0bNDMGdO" width="690" height="422" data-dominant-color="3B3D3F"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">864×529 16.4 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>I defined the function to create the custom Dataset as follow:</p>
<pre><code class="lang-auto">def get_dataset(path, dirs):

  transforms = T.Compose([
                          T.ToPILImage(),
                          T.Resize((224,224)),
                          T.RandomHorizontalFlip(),
                          T.ToTensor()
                        ])

  dataset = CustomDataset(root=path, dirs=dirs, transforms=transforms)

  return dataset
</code></pre>
<p>The name of every image is format <code>filename_label.jpg</code>.</p>
<p>I understand the error, but I don’t know how to fix it.<br>
Can you help me please?</p>

ptrblck: <p>Your <code>__getitem__</code> method is used to load and process a single sample in the default use case using the passed <code>index</code> argument.<br>
In your code snippet it seems as if you would like to load and process the entire dataset. Is this indeed the case?</p>
<p>The error is raised in:</p>
<pre><code class="lang-python">images = self.transforms(images)
</code></pre>
<p>since you are passing a <code>list</code> of image paths while a tensor or <code>PIL.Image</code> is expected.</p>

Crispolo: <p>Probably I didn’t understand how to create a custom dataset and how to use the <code>__getiem__</code>.</p>
<p>With your hit I rewrote the code:</p>
<pre><code class="lang-auto">class CustomDataset(Dataset):
  def __init__(self, root, dirs=None, transform=None):
    self.status = True

    self.root_dir = root
    self.sel_dirs = dirs
    self.transform = transform

    self.images = []
    self.labels = []
    self.__initItem()

  def getStatus(self):
    return self.status

  def __len__(self):
    if self.sel_dirs == None:
      return 0

    if len(self.images) == len(self.labels):
      return len(self.images)
    else:
      return 0

  def getImagesLabels(self):
    for (image, label) in zip(self.images, self.labels):
      print(f'{image} - {label}')

  def __checkDir(self, dirPath):
    images = []
    labels = []

    for f in os.listdir(dirPath):
      absDirPath = os.path.join(dirPath, f)

      if os.path.isdir(absDirPath):
        tmpImgs, tmpLabels = self.__checkDir(absDirPath)
        images.append(tmpImgs)
        labels.append(tmpLabels)
      elif os.path.isfile(absDirPath):
        images.append(absDirPath)
        fileName = os.path.splitext(os.path.split(absDirPath)[1])[0]
        labels.append(fileName.split('_')[-1])
        
    return (images, labels)

  def __initItem(self):
    images = []
    labels = []

    if self.sel_dirs == None:
      print(f'No selected directories')
      self.status = False
      return

    for dir in self.sel_dirs:
      dirPath = os.path.join(self.root_dir, dir)
      print('path: ' + dirPath)

      if os.path.isdir(dirPath):
        tmpImgs, tmpLabels = self.__checkDir(dirPath)
        images.extend(tmpImgs)
        labels.extend(tmpLabels)
      else:
        print(f'Data which is not a folder found into selected directories')
        self.status = False
        return

    if VERBOSE &gt;= 1:
      print('found ' + str(len(images)) + ' images and ' + str(len(labels)) + ' labels')

    self.images = images
    self.labels = labels

  def __getitem__(self, index):
    image = io.imread(self.images[index])
    label = torch.tensor(int(self.labels[index]))

    if self.transform:
      image = self.transform(image)
    
    return (image, label)
</code></pre>
<p>Now everything seems to work well.</p>