Repository: DAI-Lab/SteganoGAN

Issue:
Move `_get_steganogan` logic to `SteganoGAN.load`
The code that loads the pretrained models is currently inside a "private" function in the `cli.py` module.

This should be moved inside the `SteganoGAN.load` method to allow loading pretrained models directly from the class.

For this, the following changes are needed:

* Add an `architecture` argument to the `load` method.
* Make the `path` argument from the `load` method optional (`=None`).

Then, reimplement the load method so that:

* If a path is given, it is directly used to load the model from it.
* If an architecture name is given, the path to the model is built dynamically in a similar way to what we currently have in `_get_stegangan`.
* If either both arguments or none are given, an exception is raised.

Then, in the `cli.py` module, the _get_steganogan implementation can be changed pass the architecture value directly to the `load` method. 

Discussion:

Patch:
diff --git a/setup.py b/setup.py
index dbe8ad7..9b57672 100644
--- a/setup.py
+++ b/setup.py
@@ -74,7 +74,7 @@ setup(
     description="Steganography tool based on DeepLearning GANs",
     entry_points={
         'console_scripts': [
-            'steganogan=steganogan:cli.main'
+            'steganogan=steganogan.cli:main'
         ],
     },
     extras_require={
diff --git a/steganogan/__init__.py b/steganogan/__init__.py
index be09288..afe89d7 100644
--- a/steganogan/__init__.py
+++ b/steganogan/__init__.py
@@ -5,7 +5,6 @@ __author__ = """MIT Data To AI Lab"""
 __email__ = 'dailabmit@gmail.com'
 __version__ = '0.1.2-dev'
 
-from steganogan import cli
 from steganogan.models import SteganoGAN
 
-__all__ = ('SteganoGAN', 'cli')
+__all__ = ('SteganoGAN', )
diff --git a/steganogan/cli.py b/steganogan/cli.py
index 1ae4d3c..6a833fd 100644
--- a/steganogan/cli.py
+++ b/steganogan/cli.py
@@ -1,22 +1,23 @@
 # -*- coding: utf-8 -*-
 
-"""Top-level package for SteganoGAN."""
-
-__author__ = 'MIT Data To AI Lab'
-__email__ = 'dailabmit@gmail.com'
-__version__ = '0.1.0.dev.dev'
-
 import argparse
-import os
 
 from steganogan.models import SteganoGAN
 
 
 def _get_steganogan(args):
-    model_name = '{}.steg'.format(args.architecture)
-    pretrained_path = os.path.join(os.path.dirname(__file__), 'pretrained')
-    model_path = os.path.join(pretrained_path, model_name)
-    return SteganoGAN.load(model_path, cuda=not args.cpu, verbose=args.verbose)
+
+    steganogan_kwargs = {
+        'cuda': not args.cpu,
+        'verbose': args.verbose
+    }
+
+    if args.path:
+        steganogan_kwargs['path'] = args.path
+    else:
+        steganogan_kwargs['architecture'] = args.architecture
+
+    return SteganoGAN.load(**steganogan_kwargs)
 
 
 def _encode(args):
@@ -44,8 +45,12 @@ def _get_parser():
     # Parent Parser - Shared options
     parent = argparse.ArgumentParser(add_help=False)
     parent.add_argument('-v', '--verbose', action='store_true', help='Be verbose')
-    parent.add_argument('-a', '--architecture', default='dense', choices=('basic', 'dense'),
-                        help='Model architecture. Use the same one for both encoding and decoding')
+    group = parent.add_mutually_exclusive_group()
+    group.add_argument('-a', '--architecture', default='dense',
+                       choices={'basic', 'dense', 'residual'},
+                       help='Model architecture. Use the same one for both encoding and decoding')
+
+    group.add_argument('-p', '--path', help='Load a pretrained model from a given path.')
     parent.add_argument('--cpu', action='store_true',
                         help='Force CPU usage even if CUDA is available')
 
diff --git a/steganogan/models.py b/steganogan/models.py
index 78afdbc..f236e3e 100644
--- a/steganogan/models.py
+++ b/steganogan/models.py
@@ -341,8 +341,26 @@ class SteganoGAN(object):
         torch.save(self, path)
 
     @classmethod
-    def load(cls, path, cuda=True, verbose=False):
-        """Loads an instance of SteganoGAN from the given path."""
+    def load(cls, architecture=None, path=None, cuda=True, verbose=False):
+        """Loads an instance of SteganoGAN for the given architecture (default pretrained models)
+        or loads a pretrained model from a given path.
+
+        Args:
+            architecture(str): Name of a pretrained model to be loaded from the default models.
+            path(str): Path to custom pretrained model. *Architecture must be None.
+            cuda(bool): Force loaded model to use cuda (if available).
+            verbose(bool): Force loaded model to use or not verbose.
+        """
+
+        if architecture and not path:
+            model_name = '{}.steg'.format(architecture)
+            pretrained_path = os.path.join(os.path.dirname(__file__), 'pretrained')
+            path = os.path.join(pretrained_path, model_name)
+
+        elif (architecture is None and path is None) or (architecture and path):
+            raise ValueError(
+                'Please provide either an architecture or a path to pretrained model.')
+
         steganogan = torch.load(path, map_location='cpu')
         steganogan.verbose = verbose
 


Test patch:
diff --git a/tests/test_cli.py b/tests/test_cli.py
index db415fe..b33945e 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -1,13 +1,75 @@
 # -*- coding: utf-8 -*-
 
-import os
 from unittest.mock import MagicMock, patch
 
 from steganogan import cli
 
 
 @patch('steganogan.cli.SteganoGAN.load')
-def test__get_steganogan(mock_steganogan_load):
+def test__get_steganogan_with_path_and_architecture(mock_steganogan_load):
+    """
+    Test that:
+        * The model is called with the path and not acrhitecture.
+        * SteganoGAN.load is called with the right values
+        * The output of SteganoGAN.load is returned
+    """
+
+    # setup
+    mock_steganogan_load.return_value = 'Steganogan'
+
+    params = MagicMock(
+        architecture='dense',
+        cpu=True,
+        path='my_path/basic',
+        verbose=True,
+    )
+
+    # run
+    cli_test = cli._get_steganogan(params)
+
+    # assert
+    mock_steganogan_load.assert_called_once_with(
+        path='my_path/basic',
+        cuda=False,
+        verbose=True
+    )
+
+    assert cli_test == 'Steganogan'
+
+
+@patch('steganogan.cli.SteganoGAN.load')
+def test__get_steganogan_with_path(mock_steganogan_load):
+    """
+    Test that:
+        * The model is loaded with the path.
+        * SteganoGAN.load is called with the right values
+        * The output of SteganoGAN.load is returned
+    """
+
+    # setup
+    mock_steganogan_load.return_value = 'Steganogan'
+
+    params = MagicMock(
+        cpu=True,
+        path='my_path/basic',
+        verbose=True,
+    )
+
+    # run
+    cli_test = cli._get_steganogan(params)
+
+    # assert
+    mock_steganogan_load.assert_called_once_with(
+        path='my_path/basic',
+        cuda=False,
+        verbose=True
+    )
+
+    assert cli_test == 'Steganogan'
+
+
+@patch('steganogan.cli.SteganoGAN.load')
+def test__get_steganogan_with_architecture(mock_steganogan_load):
     """
     Test that:
         * The model path is the right one, with the right architecture
@@ -22,20 +84,15 @@ def test__get_steganogan(mock_steganogan_load):
         cpu=True,
         architecture='basic',
         verbose=True,
+        path=None
     )
 
-    model_name = '{}.steg'.format(params.architecture)
-    parent_path = os.path.dirname(os.path.dirname(__file__))
-    stega_path = os.path.join(parent_path, 'steganogan')
-    pretrained_path = os.path.join(stega_path, 'pretrained')
-    model_path = os.path.join(pretrained_path, model_name)
-
     # run
     cli_test = cli._get_steganogan(params)
 
     # assert
     mock_steganogan_load.assert_called_once_with(
-        model_path,
+        architecture='basic',
         cuda=False,
         verbose=True
     )
diff --git a/tests/test_models.py b/tests/test_models.py
index c3b9877..e4185f6 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -1067,19 +1067,20 @@ class TestSteganoGAN(TestCase):
         # assert
         mock_save.assert_called_once_with(steganogan, 'some_path')
 
+    @patch('steganogan.models.os.path.join')
     @patch('steganogan.models.torch.load')
-    def test_load(self, mock_load):
-        """Test when loading a path, we execute and update steganogan"""
+    def test_load(self, mock_load, os_path_mock):
+        """Test loading a default architecture"""
 
         # setup
         steganogan = MagicMock()
         mock_load.return_value = steganogan
 
         # run
-        result = models.SteganoGAN.load('some_path', cuda=False, verbose=False)
+        result = models.SteganoGAN.load('some_architecture', cuda=False, verbose=False)
 
         # assert
-        mock_load.assert_called_once_with('some_path', map_location='cpu')
+        mock_load.assert_called_once_with(os_path_mock.return_value, map_location='cpu')
 
         assert not steganogan.verbose
 
@@ -1091,8 +1092,9 @@ class TestSteganoGAN(TestCase):
 
         assert result == steganogan
 
+    @patch('steganogan.models.os.path.join')
     @patch('steganogan.models.torch.load')
-    def test_load_cuda_verbose_true(self, mock_load):
+    def test_load_cuda_verbose_true(self, mock_load, os_path_mock):
         """Test loading a path, with cuda and verbose True"""
 
         # setup
@@ -1100,10 +1102,10 @@ class TestSteganoGAN(TestCase):
         mock_load.return_value = steganogan
 
         # run
-        result = models.SteganoGAN.load('some_path', cuda=True, verbose=True)
+        result = models.SteganoGAN.load(architecture='some_architecture', cuda=True, verbose=True)
 
         # assert
-        mock_load.assert_called_once_with('some_path', map_location='cpu')
+        mock_load.assert_called_once_with(os_path_mock.return_value, map_location='cpu')
 
         assert steganogan.verbose
 
@@ -1114,3 +1116,26 @@ class TestSteganoGAN(TestCase):
         steganogan.set_device.assert_called_once_with(True)
 
         assert result == steganogan
+
+    @patch('steganogan.models.torch.load')
+    def test_load_path_no_architecture(self, mock_load):
+        """Test loading a model when passing a path and architecture is None"""
+
+        # setup
+        steganogan = MagicMock()
+        mock_load.return_value = steganogan
+
+        # run
+        result = models.SteganoGAN.load(
+            architecture=None, path='some_path', cuda=True, verbose=True)
+
+        # assert
+        mock_load.assert_called_once_with('some_path', map_location='cpu')
+
+        assert result == steganogan
+
+    def test_load_path_and_architecture(self):
+
+        # run / assert
+        with self.assertRaises(ValueError):
+            models.SteganoGAN.load(architecture='some_arch', path='some_path')