Repository: lightly-ai/lightly

Issue:
_lightly_cli is missing it's return statement
as a side effect `train_model_and_embed_images` is not working as expected.


Discussion:
@opassos Thanks for pointing out this bug!
I could reproduce it by changing our unittest in
https://github.com/lightly-ai/lightly/blob/a925c61bc1c9ecf1fbd3c68acf402b6dc0deb01c/tests/core/test_Core.py#L57-L58
to 
```
embeddings, labels, filenames = train_model_and_embed_images(
            input_dir=dataset_dir, trainer=trainer)
```

Which is causing it to fail with `TypeError: cannot unpack non-iterable NoneType object`.

@opassos Is this the bug you encountered?

Resolving patch:
diff --git a/lightly/cli/embed_cli.py b/lightly/cli/embed_cli.py
--- a/lightly/cli/embed_cli.py
+++ b/lightly/cli/embed_cli.py
@@ -9,8 +9,10 @@
 # All Rights Reserved
 
 import os
+from typing import Union, Tuple, List
 
 import hydra
+import numpy as np
 import torch
 import torchvision
 from torch.utils.hipify.hipify_python import bcolors
@@ -24,7 +26,21 @@
 from lightly.cli._helpers import cpu_count
 
 
-def _embed_cli(cfg, is_cli_call=True):
+def _embed_cli(cfg, is_cli_call=True) -> \
+    Union[
+        Tuple[np.ndarray, List[int], List[str]],
+        str
+    ]:
+    """ See embed_cli() for usage documentation
+
+        is_cli_call:
+            If True:
+                Saves the embeddings as file and returns the filepath.
+            If False:
+                Returns the embeddings, labels, filenames as tuple.
+                Embeddings are of shape (n_samples, embedding_size)
+                len(labels) = len(filenames) = n_samples
+    """
     input_dir = cfg['input_dir']
     if input_dir and is_cli_call:
         input_dir = fix_input_path(input_dir)
@@ -79,7 +95,7 @@ def _embed_cli(cfg, is_cli_call=True):
 
 
 @hydra.main(config_path='config', config_name='config')
-def embed_cli(cfg):
+def embed_cli(cfg) -> str:
     """Embed images from the command-line.
 
     Args:
@@ -95,6 +111,9 @@ def embed_cli(cfg):
             Path to the checkpoint of a pretrained model. If left
             empty, a pretrained model by lightly is used.
 
+    Returns:
+        The path to the created embeddings file.
+
     Examples:
         >>> # embed images with default settings and a lightly model
         >>> lightly-embed input_dir=data/
diff --git a/lightly/core.py b/lightly/core.py
--- a/lightly/core.py
+++ b/lightly/core.py
@@ -2,6 +2,9 @@
 
 # Copyright (c) 2020. Lightly AG and its affiliates.
 # All Rights Reserved
+from typing import Tuple, List
+
+import numpy as np
 
 from lightly.cli.train_cli import _train_cli
 from lightly.cli.embed_cli import _embed_cli
@@ -73,12 +76,15 @@ def _add_kwargs(cfg, kwargs):
     return cfg
 
 
-def train_model_and_embed_images(config_path: str = None, **kwargs):
+def train_model_and_embed_images(config_path: str = None, **kwargs) -> Tuple[
+    np.ndarray, List[int], List[str]
+]:
     """Train a self-supervised model and use it to embed images.
 
-    Calls the same function as lightly-magic. All arguments passed to
-    lightly-magic can also be passed to this function (see below for an
-    example).
+    First trains a modle using the _train_cli(),
+    then embeds with the _embed_cli().
+    All arguments passed to the CLI functions
+    can also be passed to this function (see below for an example).
 
     Args:
         config_path:
@@ -88,6 +94,8 @@ def train_model_and_embed_images(config_path: str = None, **kwargs):
 
     Returns:
         Embeddings, labels, and filenames of the images.
+        Embeddings are of shape (n_samples, embedding_size)
+        len(labels) = len(filenames) = n_samples
 
     Examples:
         >>> import lightly
@@ -105,14 +113,16 @@ def train_model_and_embed_images(config_path: str = None, **kwargs):
         >>> my_trainer = {max_epochs: 10}
         >>> embeddings, _, _ = lightly.train_model_and_embed_images(
         >>>     input_dir='path/to/data', trainer=my_trainer)
-        >>> # the command above is equivalent to:
-        >>> # lightly-magic input_dir='path/to/data' trainer.max_epochs=10
 
     """
     config_path = _get_config_path(config_path)
     config_args = _load_config_file(config_path)
     config_args = _add_kwargs(config_args, kwargs)
-    return _lightly_cli(config_args, is_cli_call=False)
+
+    checkpoint = _train_cli(config_args, is_cli_call=False)
+    config_args['checkpoint'] = checkpoint
+    embeddings, labels, filenames = _embed_cli(config_args, is_cli_call=False)
+    return embeddings, labels, filenames
 
 
 def train_embedding_model(config_path: str = None, **kwargs):


Test patch:
diff --git a/tests/cli/test_cli_magic.py b/tests/cli/test_cli_magic.py
--- a/tests/cli/test_cli_magic.py
+++ b/tests/cli/test_cli_magic.py
@@ -48,13 +48,29 @@ def parse_cli_string(self, cli_words: str):
         for key, value in zip(dict_keys, dict_values):
             value = value.strip('\"')
             value = value.strip('\'')
-            self.cfg[key] = value
+            try:
+                value = int(value)
+            except ValueError:
+                pass
+
+            key_subparts = key.split('.')
+            if len(key_subparts) == 1:
+                self.cfg[key] = value
+            elif len(key_subparts) == 2:
+                self.cfg[key_subparts[0]][key_subparts[1]] = value
+            else:
+                raise ValueError(
+                    f'Keys with more than 2 subparts are not supported,'
+                    f'but you entered {key}.'
+                )
+
 
     def test_parse_cli_string(self):
-        cli_string = "lightly-magic dataset_id='XYZ' upload='thumbnails'"
+        cli_string = "lightly-magic dataset_id='XYZ' upload='thumbnails' trainer.max_epochs=3"
         self.parse_cli_string(cli_string)
-        assert self.cfg["dataset_id"] == 'XYZ'
-        assert self.cfg["upload"] == 'thumbnails'
+        self.assertEqual(self.cfg["dataset_id"], 'XYZ')
+        self.assertEqual(self.cfg["upload"], 'thumbnails')
+        self.assertEqual(self.cfg['trainer']['max_epochs'], 3)
 
     def test_magic_new_dataset_name(self):
         MockedApiWorkflowClient.n_dims_embeddings_on_server = 32
@@ -68,6 +84,12 @@ def test_magic_new_dataset_id(self):
         self.parse_cli_string(cli_string)
         lightly.cli.lightly_cli(self.cfg)
 
+    def test_magic_without_upload_with_trainer(self):
+        MockedApiWorkflowClient.n_dims_embeddings_on_server = 32
+        cli_string = "lightly-magic trainer.max_epochs=1"
+        self.parse_cli_string(cli_string)
+        lightly.cli.lightly_cli(self.cfg)
+
     def tearDown(self) -> None:
         for filename in ["embeddings.csv", "embeddings_sorted.csv"]:
             try:
diff --git a/tests/core/test_Core.py b/tests/core/test_Core.py
--- a/tests/core/test_Core.py
+++ b/tests/core/test_Core.py
@@ -2,6 +2,8 @@
 import os
 import re
 import shutil
+
+import numpy as np
 import torchvision
 import tempfile
 import pytest
@@ -36,13 +38,15 @@ def create_dataset(self, n_subfolders=5, n_samples_per_subfolder=20):
                 data[0].save(os.path.join(tmp_dir,
                                           folder_names[folder_idx],
                                           sample_names[sample_idx]))
+        self.dataset_dir = tmp_dir
         return tmp_dir, folder_names, sample_names
 
 
     #@pytest.mark.slow
     def test_train_and_embed(self):
-        n_subfolders = 10
-        n_samples_per_subfolder = 10
+        n_subfolders = 3
+        n_samples_per_subfolder = 3
+        n_samples = n_subfolders * n_samples_per_subfolder
 
         # embed, no overwrites
         dataset_dir, _, _ = self.create_dataset(
@@ -51,12 +55,21 @@ def test_train_and_embed(self):
         )
 
         # train, one overwrite
-        trainer = {
-            'max_epochs': 1,
-        }
-        train_model_and_embed_images(
-            input_dir=dataset_dir, trainer=trainer)
-        shutil.rmtree(dataset_dir)
+        embeddings, labels, filenames = train_model_and_embed_images(
+            input_dir=dataset_dir,
+            trainer={'max_epochs': 1},
+            loader={'num_workers': 0},
+        )
+        self.assertEqual(len(embeddings), n_samples)
+        self.assertEqual(len(labels), n_samples)
+        self.assertEqual(len(filenames), n_samples)
+        self.assertIsInstance(embeddings[0], np.ndarray)
+        self.assertIsInstance(int(labels[0]), int)  # see if casting to int works
+        self.assertIsInstance(filenames[0], str)
+
+
+    def tearDown(self) -> None:
+        shutil.rmtree(self.dataset_dir)
         pattern = '(.*)?.ckpt$'
         for root, dirs, files in os.walk(os.getcwd()):
             for file in filter(lambda x: re.match(pattern, x), files):