Repository: lightly-ai/lightly

Issue:
Add subset from filenames functionality to LightlyDataset
The `ActiveLearningAgent` returns a list of filenames for which the user will need to get predictions. Since this can be difficult and it's important, that the order of the filenames stays the same, it would be a good option to add the following functionality to the `LightlyDataset`.

```python
dataset = LightlyDataset.subset_from_filenames(input_dir, agent.unlabeled_set)
dataloader = torch.utils.data.DataLoader(dataset, shuffle=False, batch_size=1)
```

## Question A: Which kind of base dataset are supported during the creation?

1. input_dir: `str`
2. dataset: `torchvision.datasets.ImageFolder`
3. dataset: `lightly.data._helpers.DatasetFolder`
4. dataset: `lightly.data._video.VideoDataset`
5. dataset: `generic torchvision dataset` - Problem: it does not necessarily contain filenames
6. dataset: `LightlyDataset`

## Question B: How is the `SubsetDataset` represented:

1. just like the base dataset - not compatible with Option A1
2. a `LightlyDataset`
3. a `LightlyDatasetSubset`

## Question C: How is the `SubsetDataset` created:

1. Add a method to the base dataset - only compatible with option A6
2. Add a constructor to its representation - only compatible with B2 and B3


Discussion:
The following snippet could then be added to the introduction about active learning in the docs:

```python
@torch.no_grad()
def get_predictions_by_filenames_in_order(input_dir: str,
                                          agent: ActiveLearningAgent,
                                          model: torch.nn.module):
    """

    """
    transform = torchvision.transforms.ToTensor()
    dataset = LightlyDataset.subset_from_filenames(
        input_dir, agent.unlabeled_set, transform=transform
    )
    dataloader = torch.utils.data.DataLoader(dataset, shuffle=False, batch_size=1)

    predictions = []
    for img, _, filename in dataloader:
        output = model(img).detach()
        predictions.append(output.squeeze().cpu().numpy())

    return np.array(predictions)
```
An alternative solution would be to add a function like this:

```Python
dataset = LightlyDataset(input_dir)
unlabeled_set = dataset.to_subset_by_filenames(agent.unlabeled_set)
```

This function would be much easier to implement. The disadvantage is that the user will have to work on a PyTorch `Subset` instead of a `LightlyDataset`. 
Option A6, B3, C2:
We could also create a class `LightlyDatasetSubset` which is created from a `LightlyDataset` and a `list of filenames` and is a drop-in replacement for a `LightlyDataset` by possibly inheriting from it. 
This class could just overwrite the `__getitem__ `and `__len__` methods. The new `__getitem__ ` method could be easily implemented by mapping the index w.r.t. to the subset to the index w.r.t. the base dataset and then using the `__getitem__ ` method of the base dataset.
I like this idea. Let's go with that. Preferrably name it `LightlySubset`.

Resolving patch:
diff --git a/lightly/data/dataset.py b/lightly/data/dataset.py
--- a/lightly/data/dataset.py
+++ b/lightly/data/dataset.py
@@ -6,10 +6,11 @@
 import os
 import shutil
 from PIL import Image
-from typing import List, Union
+from typing import List, Union, Callable
 
 import torch.utils.data as data
 import torchvision.datasets as datasets
+from torchvision import transforms
 
 from lightly.data._helpers import _load_dataset
 from lightly.data._helpers import DatasetFolder
@@ -53,6 +54,7 @@ def _copy_image(input_dir, output_dir, filename):
     _ensure_dir(target)
     shutil.copyfile(source, target)
 
+
 def _save_image(image, output_dir, filename, fmt):
     """Saves an image in the output directory.
 
@@ -141,8 +143,8 @@ class LightlyDataset:
 
     def __init__(self,
                  input_dir: str,
-                 transform=None,
-                 index_to_filename=None):
+                 transform: transforms.Compose = None,
+                 index_to_filename: Callable[[datasets.VisionDataset, int], str] = None):
 
         # can pass input_dir=None to create an "empty" dataset
         self.input_dir = input_dir
@@ -151,7 +153,7 @@ def __init__(self,
 
         # initialize function to get filename of image
         self.index_to_filename = _get_filename_by_index
-        if index_to_filename is  not None:
+        if index_to_filename is not None:
             self.index_to_filename = index_to_filename
 
     @classmethod
@@ -205,9 +207,8 @@ def __getitem__(self, index: int):
         """
         fname = self.index_to_filename(self.dataset, index)
         sample, target = self.dataset.__getitem__(index)
-        
-        return sample, target, fname
 
+        return sample, target, fname
 
     def __len__(self):
         """Returns the length of the dataset.
diff --git a/lightly/data/lighty_subset.py b/lightly/data/lighty_subset.py
new file mode 100644
--- /dev/null
+++ b/lightly/data/lighty_subset.py
@@ -0,0 +1,74 @@
+from typing import List, Dict, Tuple
+
+from lightly.data.dataset import LightlyDataset
+
+
+class LightlySubset(LightlyDataset):
+    def __init__(self, base_dataset: LightlyDataset, filenames_subset: List[str]):
+        """Creates a subset of a LightlyDataset.
+
+        Args:
+            base_dataset:
+                The dataset to subset from.
+            filenames_subset:
+                The filenames of the samples to be part of the subset.
+        """
+        self.base_dataset = base_dataset
+        self.filenames_subset = filenames_subset
+
+        dict_base_dataset_filename_index: Dict[str, int] = dict()
+        for index in range(len(base_dataset)):
+            fname = base_dataset.index_to_filename(self.dataset, index)
+            dict_base_dataset_filename_index[fname] = index
+
+        self.mapping_subset_index_to_baseset_index = \
+            [dict_base_dataset_filename_index[filename] for filename in filenames_subset]
+
+    def __getitem__(self, index_subset: int) -> Tuple[object, object, str]:
+        """An overwrite for indexing.
+
+        Args:
+            index_subset:
+                The index of a sample w.r.t. to the subset.
+                E.g. if index_subset == 0, the sample belonging to
+                the first filename in self.filenames_subset is returned.
+
+        Returns:
+            A tuple of the sample, its target and its filename.
+
+        """
+        index_baseset = self.mapping_subset_index_to_baseset_index[index_subset]
+        sample, target, fname = self.base_dataset.__getitem__(index_baseset)
+        return sample, target, fname
+
+    def __len__(self) -> int:
+        """Overwrites the len(...) function.
+
+        Returns:
+            The number of samples in the subset.
+        """
+        return len(self.filenames_subset)
+
+    def index_to_filename(self, dataset, index_subset: int):
+        """Maps from an index of a sample to its filename.
+
+        Args:
+            dataset:
+                Unused, but specified by the overwritten
+                function of the parent class.
+            index_subset:
+                The index of the sample w.r.t. the subset.
+
+        Returns:
+            The filename of the sample.
+        """
+        fname = self.filenames_subset[index_subset]
+        return fname
+
+    @property
+    def input_dir(self):
+        return self.base_dataset.input_dir
+
+    @property
+    def dataset(self):
+        return self.base_dataset.dataset


Test patch:
diff --git a/tests/data/test_LightlyDataset.py b/tests/data/test_LightlyDataset.py
--- a/tests/data/test_LightlyDataset.py
+++ b/tests/data/test_LightlyDataset.py
@@ -15,6 +15,7 @@
     from lightly.data._video import VideoDataset
     import av
     import cv2
+
     VIDEO_DATASET_AVAILABLE = True
 except Exception:
     VIDEO_DATASET_AVAILABLE = False
@@ -27,13 +28,13 @@ def ensure_dir(self, path_to_folder: str):
 
     def setUp(self):
         self.available_dataset_names = ['cifar10',
-                                        #'cifar100',
-                                        #'cityscapes',
-                                        #'stl10',
-                                        #'voc07-seg',
-                                        #'voc12-seg',
-                                        #'voc07-det',
-                                        #'voc12-det]
+                                        # 'cifar100',
+                                        # 'cityscapes',
+                                        # 'stl10',
+                                        # 'voc07-seg',
+                                        # 'voc12-seg',
+                                        # 'voc07-det',
+                                        # 'voc12-det]
                                         ]
 
     def create_dataset(self, n_subfolders=5, n_samples_per_subfolder=20):
@@ -63,7 +64,7 @@ def create_video_dataset(self, n_videos=5, n_frames_per_video=10, w=32, h=32, c=
 
         self.n_videos = n_videos
         self.n_frames_per_video = n_frames_per_video
-    
+
         self.input_dir = tempfile.mkdtemp()
         self.ensure_dir(self.input_dir)
         self.frames = (np.random.randn(n_frames_per_video, w, h, c) * 255).astype(np.uint8)
@@ -119,7 +120,6 @@ def test_create_lightly_dataset_from_folder_nosubdir(self):
         tmp_dir = tempfile.mkdtemp()
         sample_names = [f'img_{i}.jpg' for i in range(n_tot)]
         for sample_idx in range(n_tot):
-
             data = dataset[sample_idx]
             path = os.path.join(tmp_dir, sample_names[sample_idx])
             data[0].save(path)
@@ -150,7 +150,7 @@ def test_check_images(self):
             path = os.path.join(tmp_dir, sample_name)
             data[0].save(path)
 
-        corrupt_sample_names = [f'img_{i}.jpg' for i in range(n_healthy,n_healthy+n_corrupt)]
+        corrupt_sample_names = [f'img_{i}.jpg' for i in range(n_healthy, n_healthy + n_corrupt)]
         for sample_name in corrupt_sample_names:
             path = os.path.join(tmp_dir, sample_name)
             with open(path, 'a') as f:
@@ -158,8 +158,8 @@ def test_check_images(self):
 
         # tests
         healthy_images, corrupt_images = check_images(tmp_dir)
-        assert(len(healthy_images) == n_healthy)
-        assert(len(corrupt_images) == n_corrupt)
+        assert (len(healthy_images) == n_healthy)
+        assert (len(corrupt_images) == n_corrupt)
 
     def test_not_existing_folder_dataset(self):
         with self.assertRaises(ValueError):
@@ -173,6 +173,8 @@ def test_from_torch_dataset(self):
         self.assertEqual(len(_dataset), len(dataset))
         self.assertEqual(len(dataset.get_filenames()), len(dataset))
 
+
+
     def test_video_dataset(self):
 
         if not VIDEO_DATASET_AVAILABLE:
@@ -189,10 +191,10 @@ def test_video_dataset(self):
                 dataset = LightlyDataset(input_dir=tmp_dir)
 
             warnings.warn(
-                'Did not test video dataset because of missing requirements')        
+                'Did not test video dataset because of missing requirements')
             shutil.rmtree(tmp_dir)
             return
-        
+
         self.create_video_dataset()
         dataset = LightlyDataset(input_dir=self.input_dir)
 
@@ -200,10 +202,11 @@ def test_video_dataset(self):
         dataset.dump(out_dir)
         self.assertEqual(len(os.listdir(out_dir)), len(dataset))
 
-    def test_transform_setter(self):
-        
-        tmp_dir, _, _ = self.create_dataset()
-        dataset = LightlyDataset(input_dir=tmp_dir)
+    def test_transform_setter(self, dataset: LightlyDataset = None):
+
+        if dataset is None:
+            tmp_dir, _, _ = self.create_dataset()
+            dataset = LightlyDataset(input_dir=tmp_dir)
         # the transform of both datasets should be None
         self.assertIsNone(dataset.transform)
         self.assertIsNone(dataset.dataset.transform)
diff --git a/tests/data/test_LightlySubset.py b/tests/data/test_LightlySubset.py
new file mode 100644
--- /dev/null
+++ b/tests/data/test_LightlySubset.py
@@ -0,0 +1,51 @@
+import os
+import tempfile
+import random
+from typing import Tuple, List
+
+import torchvision
+
+from lightly.data.dataset import LightlyDataset
+from lightly.data.lighty_subset import LightlySubset
+
+from tests.data.test_LightlyDataset import TestLightlyDataset
+
+
+class TestLightlySubset(TestLightlyDataset):
+    def setUp(self) -> None:
+        tmp_dir, folder_names, sample_names = self.create_dataset(n_subfolders=5, n_samples_per_subfolder=5)
+        self.input_dir = tmp_dir
+        
+    def create_subset(self) -> Tuple[LightlySubset, List[str]]:
+        base_dataset = LightlyDataset(input_dir=self.input_dir)
+        filenames_base_dataset = base_dataset.get_filenames()
+
+        no_samples_subset = int(len(filenames_base_dataset) * 0.5)
+        filenames_subset = random.sample(filenames_base_dataset, no_samples_subset)
+
+        subset = LightlySubset(base_dataset=base_dataset, filenames_subset=filenames_subset)
+        return subset, filenames_subset
+
+    def test_create_lightly_subset(self):
+        subset, filenames_subset = self.create_subset()
+        
+        assert subset.get_filenames() == filenames_subset
+        for index_subset, filename_subset in enumerate(filenames_subset):
+            sample, target, fname = subset.__getitem__(index_subset)
+            assert filename_subset == fname
+            
+    def test_lightly_subset_transform(self):
+        subset, filenames_subset = self.create_subset()
+        self.test_transform_setter(dataset=subset)
+
+    def test_lightly_subset_dump(self):
+        subset, filenames_subset = self.create_subset()
+        dataset = subset
+
+        out_dir = tempfile.mkdtemp()
+        dataset.dump(out_dir)
+
+        files_output_dir = LightlyDataset(input_dir=out_dir).get_filenames()
+        assert set(files_output_dir) == set(dataset.get_filenames())
+
+        
\ No newline at end of file