Add ImageFolder dataset

bob/ip/binseg/configs/datasets/imagefolder.py

+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+from bob.ip.binseg.data.transforms import *
+from bob.ip.binseg.data.imagefolder import ImageFolder
+
+#### Config ####
+
+# add your transforms below
+transforms = Compose([  
+                        CenterCrop((544,544))
+                        ,RandomHFlip()
+                        ,RandomVFlip()
+                        ,RandomRotation()
+                        ,ColorJitter()
+                        ,ToTensor()
+                    ])
+
+# PyTorch dataset
+path = '/path/to/dataset'
+dataset = ImageFolder(path,transform=transforms)

bob/ip/binseg/data/imagefolder.py

+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+from torch.utils.data import Dataset
+from pathlib import Path
+import numpy as np
+from PIL import Image
+import torch
+import torchvision.transforms.functional as VF
+
+def get_file_lists(data_path):
+    data_path = Path(data_path)
+    
+    image_path = data_path.joinpath('images')
+    image_file_names = np.array(sorted(list(image_path.glob('*'))))
+
+    gt_path = data_path.joinpath('gt')
+    gt_file_names = np.array(sorted(list(gt_path.glob('*'))))
+    return image_file_names, gt_file_names
+
+class ImageFolder(Dataset):
+    """
+    Required directory tree structure for images and ground-truth (gt):
+    root
+       |- images
+       |- gt 
+    
+    
+    Parameters
+    ----------
+    path : str
+        full path to root of dataset
+    
+    Returns
+    -------
+    [type]
+        [description]
+    """
+    def __init__(self, path, transform = None):
+        self.transform = transform
+        self.img_file_list, self.gt_file_list = get_file_lists(path)
+
+    def __len__(self):
+        """
+        Returns
+        -------
+        int
+            size of the dataset
+        """
+        return len(self.img_file_list)
+    
+    def __getitem__(self,index):
+        """
+        Parameters
+        ----------
+        index : int
+        
+        Returns
+        -------
+        list
+            dataitem [img_name, img, gt, mask]
+        """
+        img_path = self.img_file_list[index]
+        img_name = img_path.name
+        img = Image.open(img_path).convert(mode='RGB')
+    
+        gt_path = self.gt_file_list[index]
+        gt = Image.open(gt_path).convert(mode='1', dither=None)
+        sample = [img, gt]
+        
+        if self.transform :
+            sample = self.transform(*sample)
+        
+        sample.insert(0,img_name)
+        
+        return sample

doc/configs.rst

@@ -8,6 +8,13 @@ Configs
 Dataset Configs
 ===============
 
+.. _bob.ip.binseg.configs.datasets.imagefolder:
+
+ImageFolder
+----------------
+.. literalinclude:: ../bob/ip/binseg/configs/datasets/imagefolder.py
+
+
 .. _bob.ip.binseg.configs.datasets.chasedb1:
 
 CHASEDB1

doc/datasets.rst

@@ -33,9 +33,16 @@ Supported Datasets
 Add-on: Folder-based Dataset
 ============================
 
-For quick experimentation we also provide a PyTorch class that works with a 
-dataset folder structure:
+For quick experimentation we also provide a PyTorch class that works with the following 
+dataset folder structure for images and ground-truth (gt):
 
+.. code-block:: bash
 
+    root
+       |- images
+       |- gt 
+
+In the dataset config :ref:`bob.ip.binseg.configs.datasets.imagefolder` the full path of the dataset has to be amended. Training can then for example be started with
+``bob binseg train M2UNet IMAGEFOLDER -b 4 -d cuda -o /my/output/path -vv``
 
 .. include:: links.rst

requirements.txt

 setuptools
 numpy
 bob.extension
-bob.db.drive
-bob.db.stare
-bob.db.chasedb1
-bob.db.hrf
-bob.db.drionsdb
-bob.db.rimoner3
-bob.db.drishtigs1
-bob.db.refuge
-bob.db.iostar
 torch
 torchvision
 pandas

setup.py

@@ -65,6 +65,7 @@ setup(
           'M2UNetSSL = bob.ip.binseg.configs.models.m2unetssl',
           'UNet = bob.ip.binseg.configs.models.unet',
           'ResUNet = bob.ip.binseg.configs.models.resunet',
+          'IMAGEFOLDER = bob.ip.binseg.configs.datasets.imagefolder',
           'CHASEDB1 = bob.ip.binseg.configs.datasets.chasedb1',
           'CHASEDB1TEST = bob.ip.binseg.configs.datasets.chasedb1test',
           'COVD-DRIVE = bob.ip.binseg.configs.datasets.starechasedb1iostarhrf544',