diff --git a/README.md b/README.md
index 6d0e44c..685e9d9 100644
--- a/README.md
+++ b/README.md
@@ -68,7 +68,7 @@ Secret Message
 ### Hide the message with the Sieve of Eratosthenes
 
 ```bash
-$ stegano-lsb-set hide -i ./tests/sample-files/Lenna.png -m 'Secret Message' --generator eratosthenes -o Lena2.png
+$ stegano-lsb hide -i ./tests/sample-files/Lenna.png -m 'Secret Message' --generator eratosthenes -o Lena2.png
 ```
 
 The message will be scattered in the picture, following a set described by the
diff --git a/bin/__init__.py b/bin/__init__.py
index 8b13789..e69de29 100644
--- a/bin/__init__.py
+++ b/bin/__init__.py
@@ -1 +0,0 @@
-
diff --git a/bin/lsb.py b/bin/lsb.py
index 1e8d9e7..7dd342c 100755
--- a/bin/lsb.py
+++ b/bin/lsb.py
@@ -20,20 +20,38 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>
 
 __author__ = "Cedric Bonhomme"
-__version__ = "$Revision: 0.8 $"
-__date__ = "$Date: 2016/08/04 $"
-__revision__ = "$Date: 2019/06/01 $"
+__version__ = "$Revision: 0.7 $"
+__date__ = "$Date: 2016/03/18 $"
+__revision__ = "$Date: 2019/06/04 $"
 __license__ = "GPLv3"
 
-import argparse
+import inspect
+import crayons
 
 try:
     from stegano import lsb
+    from stegano.lsb import generators
 except Exception:
-    print("Install Stegano: pipx install Stegano")
+    print("Install stegano: pipx install Stegano")
 
 from stegano import tools
 
+import argparse
+
+
+class ValidateGenerator(argparse.Action):
+    def __call__(self, parser, args, values, option_string=None):
+        valid_generators = [
+            generator[0]
+            for generator in inspect.getmembers(generators, inspect.isfunction)
+        ]
+        # Verify that the generator is valid
+        generator = values[0]
+        if generator not in valid_generators:
+            raise ValueError("Unknown generator: %s" % generator)
+        # Set the generator_function arg of the parser
+        setattr(args, self.dest, values)
+
 
 def main():
     parser = argparse.ArgumentParser(prog="stegano-lsb")
@@ -59,6 +77,22 @@ def main():
         + " UTF-8 (default) or UTF-32LE.",
     )
 
+    # Generator
+    parser_hide.add_argument(
+        "-g",
+        "--generator",
+        dest="generator_function",
+        action=ValidateGenerator,
+        nargs="*",
+        required=False,
+        help="Generator (with optional arguments)",
+    )
+
+    # Shift the message to hide
+    parser_hide.add_argument(
+        "-s", "--shift", dest="shift", default=0, help="Shift for the generator"
+    )
+
     group_secret = parser_hide.add_mutually_exclusive_group(required=True)
     # Non binary secret message to hide
     group_secret.add_argument(
@@ -78,11 +112,6 @@ def main():
         help="Output image containing the secret.",
     )
 
-    # Shift the message to hide
-    parser_hide.add_argument(
-        "-s", "--shift", dest="shift", default=0, help="Shift for the message to hide"
-    )
-
     # Subparser: Reveal
     parser_reveal = subparsers.add_parser("reveal", help="reveal help")
     parser_reveal.add_argument(
@@ -101,26 +130,69 @@ def main():
         help="Specify the encoding of the message to reveal."
         + " UTF-8 (default) or UTF-32LE.",
     )
+
+    # Generator
+    parser_reveal.add_argument(
+        "-g",
+        "--generator",
+        dest="generator_function",
+        action=ValidateGenerator,
+        nargs="*",
+        required=False,
+        help="Generator (with optional arguments)",
+    )
+
+    # Shift the message to reveal
+    parser_reveal.add_argument(
+        "-s", "--shift", dest="shift", default=0, help="Shift for the generator"
+    )
     parser_reveal.add_argument(
         "-o",
         dest="secret_binary",
         help="Output for the binary secret (Text or any binary file).",
     )
-    # Shift the message to reveal
-    parser_reveal.add_argument(
-        "-s", "--shift", dest="shift", default=0, help="Shift for the reveal"
+
+    # Subparser: List generators
+    parser_list_generators = subparsers.add_parser(
+        "list-generators", help="list-generators help"
     )
 
     arguments = parser.parse_args()
 
+    if arguments.command != "list-generators":
+        if not arguments.generator_function:
+            generator = None
+        else:
+            try:
+                if arguments.generator_function[0] == "LFSR":
+                    # Compute the size of the image for use by the LFSR generator if needed
+                    tmp = tools.open_image(arguments.input_image_file)
+                    size = tmp.width * tmp.height
+                    tmp.close()
+                    arguments.generator_function.append(size)
+                if len(arguments.generator_function) > 1:
+                    generator = getattr(generators, arguments.generator_function[0])(
+                        *[int(e) for e in arguments.generator_function[1:]]
+                    )
+                else:
+                    generator = getattr(generators, arguments.generator_function[0])()
+
+            except AttributeError:
+                print("Unknown generator: {}".format(arguments.generator_function))
+                exit(1)
+
     if arguments.command == "hide":
         if arguments.secret_message is not None:
             secret = arguments.secret_message
-        elif arguments.secret_file is not None:
+        elif arguments.secret_file != "":
             secret = tools.binary2base64(arguments.secret_file)
 
         img_encoded = lsb.hide(
-            arguments.input_image_file, secret, arguments.encoding, int(arguments.shift)
+            image=arguments.input_image_file,
+            message=secret,
+            generator=generator,
+            shift=int(arguments.shift),
+            encoding=arguments.encoding,
         )
         try:
             img_encoded.save(arguments.output_image_file)
@@ -129,12 +201,27 @@ def main():
             print(e)
 
     elif arguments.command == "reveal":
-        secret = lsb.reveal(
-            arguments.input_image_file, arguments.encoding, int(arguments.shift)
-        )
+        try:
+            secret = lsb.reveal(
+                encoded_image=arguments.input_image_file,
+                generator=generator,
+                shift=int(arguments.shift),
+                encoding=arguments.encoding,
+            )
+        except IndexError:
+            print("Impossible to detect message.")
+            exit(0)
         if arguments.secret_binary is not None:
             data = tools.base642binary(secret)
             with open(arguments.secret_binary, "wb") as f:
                 f.write(data)
         else:
             print(secret)
+
+    elif arguments.command == "list-generators":
+        all_generators = inspect.getmembers(generators, inspect.isfunction)
+        for generator in all_generators:
+            print("Generator id:")
+            print("    {}".format(crayons.green(generator[0], bold=True)))
+            print("Desciption:")
+            print("    {}".format(generator[1].__doc__))
diff --git a/bin/lsbset.py b/bin/lsbset.py
deleted file mode 100755
index ffba1a0..0000000
--- a/bin/lsbset.py
+++ /dev/null
@@ -1,218 +0,0 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-
-# Stegano - Stegano is a pure Python steganography module.
-# Copyright (C) 2010-2022 Cédric Bonhomme - https://www.cedricbonhomme.org
-#
-# For more information : https://github.com/cedricbonhomme/Stegano
-#
-# This program is free software: you can redistribute it and/or modify
-# it under the terms of the GNU General Public License as published by
-# the Free Software Foundation, either version 3 of the License, or
-# (at your option) any later version.
-#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-# GNU General Public License for more details.
-#
-# You should have received a copy of the GNU General Public License
-# along with this program.  If not, see <http://www.gnu.org/licenses/>
-
-__author__ = "Cedric Bonhomme"
-__version__ = "$Revision: 0.7 $"
-__date__ = "$Date: 2016/03/18 $"
-__revision__ = "$Date: 2019/06/04 $"
-__license__ = "GPLv3"
-
-import inspect
-import crayons
-
-try:
-    from stegano import lsbset
-    from stegano.lsbset import generators
-except Exception:
-    print("Install stegano: pipx install Stegano")
-
-from stegano import tools
-
-import argparse
-
-
-class ValidateGenerator(argparse.Action):
-    def __call__(self, parser, args, values, option_string=None):
-        valid_generators = [
-            generator[0]
-            for generator in inspect.getmembers(generators, inspect.isfunction)
-        ]
-        # Verify that the generator is valid
-        generator = values[0]
-        if generator not in valid_generators:
-            raise ValueError("Unknown generator: %s" % generator)
-        # Set the generator_function arg of the parser
-        setattr(args, self.dest, values)
-
-
-def main():
-    parser = argparse.ArgumentParser(prog="stegano-lsb-set")
-    subparsers = parser.add_subparsers(help="sub-command help", dest="command")
-
-    # Subparser: Hide
-    parser_hide = subparsers.add_parser("hide", help="hide help")
-    # Original image
-    parser_hide.add_argument(
-        "-i",
-        "--input",
-        dest="input_image_file",
-        required=True,
-        help="Input image file.",
-    )
-    parser_hide.add_argument(
-        "-e",
-        "--encoding",
-        dest="encoding",
-        choices=tools.ENCODINGS.keys(),
-        default="UTF-8",
-        help="Specify the encoding of the message to hide."
-        + " UTF-8 (default) or UTF-32LE.",
-    )
-
-    # Generator
-    parser_hide.add_argument(
-        "-g",
-        "--generator",
-        dest="generator_function",
-        action=ValidateGenerator,
-        nargs="*",
-        required=True,
-        help="Generator (with optional arguments)",
-    )
-    parser_hide.add_argument(
-        "-s", "--shift", dest="shift", default=0, help="Shift for the generator"
-    )
-
-    group_secret = parser_hide.add_mutually_exclusive_group(required=True)
-    # Non binary secret message to hide
-    group_secret.add_argument(
-        "-m", dest="secret_message", help="Your secret message to hide (non binary)."
-    )
-    # Binary secret message to hide
-    group_secret.add_argument(
-        "-f", dest="secret_file", help="Your secret to hide (Text or any binary file)."
-    )
-
-    # Image containing the secret
-    parser_hide.add_argument(
-        "-o",
-        "--output",
-        dest="output_image_file",
-        required=True,
-        help="Output image containing the secret.",
-    )
-
-    # Subparser: Reveal
-    parser_reveal = subparsers.add_parser("reveal", help="reveal help")
-    parser_reveal.add_argument(
-        "-i",
-        "--input",
-        dest="input_image_file",
-        required=True,
-        help="Input image file.",
-    )
-    parser_reveal.add_argument(
-        "-e",
-        "--encoding",
-        dest="encoding",
-        choices=tools.ENCODINGS.keys(),
-        default="UTF-8",
-        help="Specify the encoding of the message to reveal."
-        + " UTF-8 (default) or UTF-32LE.",
-    )
-    parser_reveal.add_argument(
-        "-g",
-        "--generator",
-        dest="generator_function",
-        action=ValidateGenerator,
-        nargs="*",
-        required=True,
-        help="Generator (with optional arguments)",
-    )
-    parser_reveal.add_argument(
-        "-s", "--shift", dest="shift", default=0, help="Shift for the generator"
-    )
-    parser_reveal.add_argument(
-        "-o",
-        dest="secret_binary",
-        help="Output for the binary secret (Text or any binary file).",
-    )
-
-    # Subparser: List generators
-    parser_list_generators = subparsers.add_parser(
-        "list-generators", help="list-generators help"
-    )
-
-    arguments = parser.parse_args()
-
-    if arguments.command != "list-generators":
-        try:
-            arguments.generator_function[0]
-        except AttributeError:
-            print("You must specify the name of a generator.")
-            parser.print_help()
-            exit(1)
-
-        try:
-            if arguments.generator_function[0] == "LFSR":
-                # Compute the size of the image for use by the LFSR generator if needed
-                tmp = tools.open_image(arguments.input_image_file)
-                size = tmp.width * tmp.height
-                tmp.close()
-                arguments.generator_function.append(size)
-            if len(arguments.generator_function) > 1:
-                generator = getattr(generators, arguments.generator_function[0])(
-                    *[int(e) for e in arguments.generator_function[1:]]
-                )
-            else:
-                generator = getattr(generators, arguments.generator_function[0])()
-
-        except AttributeError:
-            print("Unknown generator: {}".format(arguments.generator_function))
-            exit(1)
-
-    if arguments.command == "hide":
-        if arguments.secret_message is not None:
-            secret = arguments.secret_message
-        elif arguments.secret_file != "":
-            secret = tools.binary2base64(arguments.secret_file)
-
-        img_encoded = lsbset.hide(
-            arguments.input_image_file, secret, generator, int(arguments.shift)
-        )
-        try:
-            img_encoded.save(arguments.output_image_file)
-        except Exception as e:
-            # If hide() returns an error (Too long message).
-            print(e)
-
-    elif arguments.command == "reveal":
-        try:
-            secret = lsbset.reveal(
-                arguments.input_image_file, generator, int(arguments.shift)
-            )
-        except IndexError:
-            print("Impossible to detect message.")
-            exit(0)
-        if arguments.secret_binary is not None:
-            data = tools.base642binary(secret)
-            with open(arguments.secret_binary, "wb") as f:
-                f.write(data)
-        else:
-            print(secret)
-
-    elif arguments.command == "list-generators":
-        all_generators = inspect.getmembers(generators, inspect.isfunction)
-        for generator in all_generators:
-            print("Generator id:")
-            print("    {}".format(crayons.green(generator[0], bold=True)))
-            print("Desciption:")
-            print("    {}".format(generator[1].__doc__))
diff --git a/docs/module.rst b/docs/module.rst
index bfe01c1..13040a8 100644
--- a/docs/module.rst
+++ b/docs/module.rst
@@ -28,26 +28,31 @@ Sets are used in order to select the pixels where the message will be hidden.
 
     Python 3.10.0 (default, Oct 17 2021, 09:02:57) [GCC 11.2.0] on linux
     Type "help", "copyright", "credits" or "license" for more information.
-    >>> from stegano import lsbset
-    >>> from stegano.lsbset import generators
+    >>> from stegano import lsb
+    >>> from stegano.lsb import generators
 
     # Hide a secret with the Sieve of Eratosthenes
     >>> secret_message = "Hello World!"
-    >>> secret_image = lsbset.hide("./tests/sample-files/Lenna.png",
-                                    secret_message,
-                                    generators.eratosthenes())
+    >>> secret_image = lsb.hide("./tests/sample-files/Lenna.png", secret_message, generators.eratosthenes())
     >>> secret_image.save("./image.png")
 
     # Try to decode with another generator
-    >>> message = lsbset.reveal("./image.png", generators.fibonacci())
+    >>> message = lsb.reveal("./image.png", generators.fibonacci())
     Traceback (most recent call last):
-      File "<stdin>", line 1, in <module>
-      File "/home/cedric/projects/Stegano/stegano/lsbset/lsbset.py", line 111, in reveal
-        for color in img_list[generated_number]:
-    IndexError: list index out of range
+    File "/Users/flavien/.local/share/virtualenvs/Stegano-sY_cwr69/bin/stegano-lsb", line 6, in <module>
+        sys.exit(main())
+    File "/Users/flavien/Perso/dev/Stegano/bin/lsb.py", line 190, in main
+        img_encoded = lsb.hide(
+    File "/Users/flavien/Perso/dev/Stegano/stegano/lsb/lsb.py", line 63, in hide
+        hider.encode_pixel((col, row))
+    File "/Users/flavien/Perso/dev/Stegano/stegano/tools.py", line 165, in encode_pixel
+        r, g, b, *a = self.encoded_image.getpixel(coordinate)
+    File "/Users/flavien/.local/share/virtualenvs/Stegano-sY_cwr69/lib/python3.10/site-packages/PIL/Image.py", line 1481, in getpixel
+        return self.im.getpixel(xy)
+    IndexError: image index out of range
 
     # Decode with Eratosthenes
-    >>> message = lsbset.reveal("./image.png", generators.eratosthenes())
+    >>> message = lsb.reveal("./image.png", generators.eratosthenes())
     >>> message
     'Hello World!'
 
diff --git a/docs/software.rst b/docs/software.rst
index 3e232f8..dfc9b95 100644
--- a/docs/software.rst
+++ b/docs/software.rst
@@ -12,52 +12,56 @@ Display help
 .. code-block:: bash
 
     $ stegano-lsb --help
-    usage: stegano-lsb [-h] {hide,reveal} ...
+    usage: stegano-lsb [-h] {hide,reveal,list-generators} ...
 
     positional arguments:
-      {hide,reveal}  sub-command help
-        hide         hide help
-        reveal       reveal help
+    {hide,reveal,list-generators}
+                            sub-command help
+        hide                hide help
+        reveal              reveal help
+        list-generators     list-generators help
 
-    optional arguments:
-      -h, --help     show this help message and exit
+    options:
+    -h, --help            show this help message and exit
 
 
 .. code-block:: bash
 
     $ stegano-lsb hide --help
-    usage: stegano-lsb hide [-h] -i INPUT_IMAGE_FILE [-e {UTF-8,UTF-32LE}]
-                        (-m SECRET_MESSAGE | -f SECRET_FILE) -o
-                        OUTPUT_IMAGE_FILE
+    usage: stegano-lsb hide [-h] -i INPUT_IMAGE_FILE [-e {UTF-8,UTF-32LE}] [-g [GENERATOR_FUNCTION ...]] [-s SHIFT] (-m SECRET_MESSAGE | -f SECRET_FILE) -o OUTPUT_IMAGE_FILE
 
-    optional arguments:
-      -h, --help            show this help message and exit
-      -i INPUT_IMAGE_FILE, --input INPUT_IMAGE_FILE
+    options:
+    -h, --help            show this help message and exit
+    -i INPUT_IMAGE_FILE, --input INPUT_IMAGE_FILE
                             Input image file.
-      -e {UTF-8,UTF-32LE}, --encoding {UTF-8,UTF-32LE}
-                            Specify the encoding of the message to hide. UTF-8
-                            (default) or UTF-32LE.
-      -m SECRET_MESSAGE     Your secret message to hide (non binary).
-      -f SECRET_FILE        Your secret to hide (Text or any binary file).
-      -o OUTPUT_IMAGE_FILE, --output OUTPUT_IMAGE_FILE
+    -e {UTF-8,UTF-32LE}, --encoding {UTF-8,UTF-32LE}
+                            Specify the encoding of the message to hide. UTF-8 (default) or UTF-32LE.
+    -g [GENERATOR_FUNCTION ...], --generator [GENERATOR_FUNCTION ...]
+                            Generator (with optional arguments)
+    -s SHIFT, --shift SHIFT
+                            Shift for the generator
+    -m SECRET_MESSAGE     Your secret message to hide (non binary).
+    -f SECRET_FILE        Your secret to hide (Text or any binary file).
+    -o OUTPUT_IMAGE_FILE, --output OUTPUT_IMAGE_FILE
                             Output image containing the secret.
 
 
 .. code-block:: bash
 
     $ stegano-lsb reveal --help
-    usage: stegano-lsb reveal [-h] -i INPUT_IMAGE_FILE [-e {UTF-8,UTF-32LE}]
-                              [-o SECRET_BINARY]
+    usage: stegano-lsb reveal [-h] -i INPUT_IMAGE_FILE [-e {UTF-8,UTF-32LE}] [-g [GENERATOR_FUNCTION ...]] [-s SHIFT] [-o SECRET_BINARY]
 
-    optional arguments:
-      -h, --help            show this help message and exit
-      -i INPUT_IMAGE_FILE, --input INPUT_IMAGE_FILE
+    options:
+    -h, --help            show this help message and exit
+    -i INPUT_IMAGE_FILE, --input INPUT_IMAGE_FILE
                             Input image file.
-      -e {UTF-8,UTF-32LE}, --encoding {UTF-8,UTF-32LE}
-                            Specify the encoding of the message to reveal. UTF-8
-                            (default) or UTF-32LE.
-      -o SECRET_BINARY      Output for the binary secret (Text or any binary
-                            file).
+    -e {UTF-8,UTF-32LE}, --encoding {UTF-8,UTF-32LE}
+                            Specify the encoding of the message to reveal. UTF-8 (default) or UTF-32LE.
+    -g [GENERATOR_FUNCTION ...], --generator [GENERATOR_FUNCTION ...]
+                            Generator (with optional arguments)
+    -s SHIFT, --shift SHIFT
+                            Shift for the generator
+    -o SECRET_BINARY      Output for the binary secret (Text or any binary file).
 
 
 Hide and reveal a text message
@@ -92,63 +96,40 @@ Hide and reveal a binary file
 
 
 
-
-
-
-The command ``stegano-lsb-set``
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
 Sets are used in order to select the pixels where the message will be hidden.
 
-Hide and reveal a text message
-------------------------------
+Hide and reveal a text message with set
+---------------------------------------
 
 .. code-block:: bash
 
     # Hide the message with the Sieve of Eratosthenes
-    $ stegano-lsb-set hide -i ./tests/sample-files/Montenach.png --generator eratosthenes -m 'Joyeux Noël!' -o ./surprise.png
+    $ stegano-lsb hide -i ./tests/sample-files/Montenach.png --generator eratosthenes -m 'Joyeux Noël!' -o ./surprise.png
 
     # Try to reveal with Mersenne numbers
-    $ stegano-lsb-set reveal --generator mersenne -i ./surprise.png
+    $ stegano-lsb reveal --generator mersenne -i ./surprise.png
 
     # Try to reveal with fermat numbers
-    $ stegano-lsb-set reveal --generator fermat -i ./surprise.png
+    $ stegano-lsb reveal --generator fermat -i ./surprise.png
 
     # Try to reveal with carmichael numbers
-    $ stegano-lsb-set reveal --generator carmichael -i ./surprise.png
+    $ stegano-lsb reveal --generator carmichael -i ./surprise.png
 
     # Try to reveal with Sieve of Eratosthenes
-    $ stegano-lsb-set reveal --generator eratosthenes -i ./surprise.png
-
-An other example:
-
-.. code-block:: bash
-
-    # Hide the message - LSB with a set defined by the identity function (f(x) = x).
-    stegano-lsb-set hide -i ./tests/sample-files/Montenach.png --generator identity -m 'I like steganography.' -o ./enc-identity.png
-
-    # Hide the message - LSB only.
-    stegano-lsb hide -i ./tests/sample-files/Montenach.png -m 'I like steganography.' -o ./enc.png
-
-    # Check if the two generated files are the same.
-    sha1sum ./enc-identity.png ./enc.png
-
-    # The output of lsb is given to lsb-set.
-    stegano-lsb-set reveal -i ./enc.png --generator identity
-
-    # The output of lsb-set is given to lsb.
-    stegano-lsb reveal -i ./enc-identity.png
+    $ stegano-lsb reveal --generator eratosthenes -i ./surprise.png
 
 
 Sometimes it can be useful to skip the first values of a set. For example if you want
 to hide several messages or because due to the selected generator
 (Fibonacci starts with 0, 1, 1, etc.). Or maybe you just want to add more complexity.
-In this case, simply use the optional arguments ``--shift``:
+In this case, simply use the optional arguments ``--shift`` or ``-s``:
 
 
 .. code-block:: bash
 
-    stegano-lsb-set reveal -i ./tests/sample-files/Lenna.png --generator fibonacci --shift 7
+    $ stegano-lsb hide -i ./tests/sample-files/Lenna.png -m 'Shifted secret message' -o ~/Lenna1.png --shift 7
+    $ stegano-lsb reveal -i ~/Lenna1.png --shift 7
+    Shifted secret message
 
 
 List all available generators
@@ -156,7 +137,7 @@ List all available generators
 
 .. code-block:: bash
 
-    $ stegano-lsb-set list-generators
+    $ stegano-lsb list-generators
     Generator id:
         ackermann
     Desciption:
diff --git a/docs/steganalysis.rst b/docs/steganalysis.rst
index c1f4c73..3a8b60f 100644
--- a/docs/steganalysis.rst
+++ b/docs/steganalysis.rst
@@ -7,7 +7,7 @@ Parity
 .. code-block:: bash
 
     # Hide the message  with Sieve of Eratosthenes
-    stegano-lsb-set hide -i ./tests/sample-files/20160505T130442.jpg -o ./surprise.png --generator eratosthenes -m 'Very important message.'
+    stegano-lsb hide -i ./tests/sample-files/20160505T130442.jpg -o ./surprise.png --generator eratosthenes -m 'Very important message.'
 
     # Steganalysis of the original photo
     stegano-steganalysis-parity -i ./tests/sample-files/20160505T130442.jpg -o ./surprise_st_original.png
@@ -16,4 +16,4 @@ Parity
     stegano-steganalysis-parity -i ./surprise.png -o ./surprise_st_secret.png
 
     # Reveal with Sieve of Eratosthenes
-    stegano-lsb-set reveal -i ./surprise.png --generator eratosthenes
+    stegano-lsb reveal -i ./surprise.png --generator eratosthenes
diff --git a/pyproject.toml b/pyproject.toml
index 859d42d..fc4f6ef 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -38,7 +38,6 @@ include = [
 
 [tool.poetry.scripts]
 stegano-lsb = "bin.lsb:main"
-stegano-lsb-set = "bin.lsbset:main"
 stegano-red = "bin.red:main"
 stegano-steganalysis-parity = "bin.parity:main"
 stegano-steganalysis-statistics = "bin.statistics:main"
diff --git a/stegano/__init__.py b/stegano/__init__.py
index c76943b..4bca0e9 100755
--- a/stegano/__init__.py
+++ b/stegano/__init__.py
@@ -4,6 +4,5 @@
 from . import red
 from . import exifHeader
 from . import lsb
-from . import lsbset
 
 from . import steganalysis
diff --git a/stegano/lsbset/generators.py b/stegano/lsb/generators.py
similarity index 98%
rename from stegano/lsbset/generators.py
rename to stegano/lsb/generators.py
index c450afe..ea822c3 100644
--- a/stegano/lsbset/generators.py
+++ b/stegano/lsb/generators.py
@@ -29,7 +29,7 @@ import itertools
 import cv2
 import numpy as np
 import math
-from typing import Dict, Iterator, List, Any, Union
+from typing import Dict, Iterator, List, Any
 
 
 def identity() -> Iterator[int]:
@@ -225,7 +225,7 @@ def LFSR(m: int) -> Iterator[int]:
         # Add the feedback bit
         state.insert(0, feedback)
         # Convert the registers to an int
-        out = sum([e * (2 ** i) for i, e in enumerate(state)])
+        out = sum([e * (2**i) for i, e in enumerate(state)])
         yield out
 
 
diff --git a/stegano/lsb/lsb.py b/stegano/lsb/lsb.py
old mode 100755
new mode 100644
index 7d12790..6dc670d
--- a/stegano/lsb/lsb.py
+++ b/stegano/lsb/lsb.py
@@ -20,111 +20,71 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>
 
 __author__ = "Cedric Bonhomme"
-__version__ = "$Revision: 0.4 $"
-__date__ = "$Date: 2016/08/04 $"
-__revision__ = "$Date: 2019/06/01 $"
+__version__ = "$Revision: 0.7 $"
+__date__ = "$Date: 2016/03/13 $"
+__revision__ = "$Date: 2019/05/31 $"
 __license__ = "GPLv3"
 
-from typing import IO, Union
+from typing import IO, Iterator, Union
 
+from .generators import identity
 from stegano import tools
 
 
 def hide(
-    input_image: Union[str, IO[bytes]],
+    image: Union[str, IO[bytes]],
     message: str,
-    encoding: str = "UTF-8",
+    generator: Iterator[int] = None,
     shift: int = 0,
+    encoding: str = "UTF-8",
     auto_convert_rgb: bool = False,
 ):
     """Hide a message (string) in an image with the
     LSB (Least Significant Bit) technique.
     """
-    message_length = len(message)
-    assert message_length != 0, "message length is zero"
+    hider = tools.Hider(image, message, encoding, auto_convert_rgb)
+    width = hider.encoded_image.width
 
-    img = tools.open_image(input_image)
+    if not generator:
+        generator = identity()
 
-    if img.mode not in ["RGB", "RGBA"]:
-        if not auto_convert_rgb:
-            print("The mode of the image is not RGB. Mode is {}".format(img.mode))
-            answer = input("Convert the image to RGB ? [Y / n]\n") or "Y"
-            if answer.lower() == "n":
-                raise Exception("Not a RGB image.")
-        img = img.convert("RGB")
+    while shift != 0:
+        next(generator)
+        shift -= 1
 
-    encoded = img.copy()
-    width, height = img.size
-    index = 0
+    while hider.encode_another_pixel():
+        generated_number = next(generator)
 
-    message = str(message_length) + ":" + str(message)
-    message_bits = "".join(tools.a2bits_list(message, encoding))
-    message_bits += "0" * ((3 - (len(message_bits) % 3)) % 3)
+        col = generated_number % width
+        row = int(generated_number / width)
 
-    npixels = width * height
-    len_message_bits = len(message_bits)
-    if len_message_bits > npixels * 3:
-        raise Exception(
-            "The message you want to hide is too long: {}".format(message_length)
-        )
-    for row in range(height):
-        for col in range(width):
-            if shift != 0:
-                shift -= 1
-                continue
-            if index + 3 <= len_message_bits:
+        hider.encode_pixel((col, row))
 
-                # Get the colour component.
-                pixel = img.getpixel((col, row))
-                r = pixel[0]
-                g = pixel[1]
-                b = pixel[2]
-
-                # Change the Least Significant Bit of each colour component.
-                r = tools.setlsb(r, message_bits[index])
-                g = tools.setlsb(g, message_bits[index + 1])
-                b = tools.setlsb(b, message_bits[index + 2])
-
-                # Save the new pixel
-                if img.mode == "RGBA":
-                    encoded.putpixel((col, row), (r, g, b, pixel[3]))
-                else:
-                    encoded.putpixel((col, row), (r, g, b))
-
-                index += 3
-            else:
-                img.close()
-                return encoded
+    return hider.encoded_image
 
 
-def reveal(input_image: Union[str, IO[bytes]], encoding: str = "UTF-8", shift: int = 0):
+def reveal(
+    encoded_image: Union[str, IO[bytes]],
+    generator: Iterator[int] = None,
+    shift: int = 0,
+    encoding: str = "UTF-8",
+):
     """Find a message in an image (with the LSB technique)."""
-    img = tools.open_image(input_image)
-    width, height = img.size
-    buff, count = 0, 0
-    bitab = []
-    limit = None
-    for row in range(height):
-        for col in range(width):
-            if shift != 0:
-                shift -= 1
-                continue
-            # pixel = [r, g, b] or [r,g,b,a]
-            pixel = img.getpixel((col, row))
-            if img.mode == "RGBA":
-                pixel = pixel[:3]  # ignore the alpha
-            for color in pixel:
-                buff += (color & 1) << (tools.ENCODINGS[encoding] - 1 - count)
-                count += 1
-                if count == tools.ENCODINGS[encoding]:
-                    bitab.append(chr(buff))
-                    buff, count = 0, 0
-                    if bitab[-1] == ":" and limit is None:
-                        try:
-                            limit = int("".join(bitab[:-1]))
-                        except Exception:
-                            pass
+    revealer = tools.Revealer(encoded_image, encoding)
+    width = revealer.encoded_image.width
 
-            if len(bitab) - len(str(limit)) - 1 == limit:
-                img.close()
-                return "".join(bitab)[len(str(limit)) + 1 :]
+    if not generator:
+        generator = identity()
+
+    while shift != 0:
+        next(generator)
+        shift -= 1
+
+    while True:
+        generated_number = next(generator)
+
+        col = generated_number % width
+        row = int(generated_number / width)
+
+        if revealer.decode_pixel((col, row)):
+            return revealer.secret_message
diff --git a/stegano/lsbset/__init__.py b/stegano/lsbset/__init__.py
deleted file mode 100644
index 22954f0..0000000
--- a/stegano/lsbset/__init__.py
+++ /dev/null
@@ -1,4 +0,0 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-
-from .lsbset import *
diff --git a/stegano/lsbset/lsbset.py b/stegano/lsbset/lsbset.py
deleted file mode 100644
index f6d0ee4..0000000
--- a/stegano/lsbset/lsbset.py
+++ /dev/null
@@ -1,136 +0,0 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-
-# Stegano - Stegano is a pure Python steganography module.
-# Copyright (C) 2010-2022 Cédric Bonhomme - https://www.cedricbonhomme.org
-#
-# For more information : https://git.sr.ht/~cedric/stegano
-#
-# This program is free software: you can redistribute it and/or modify
-# it under the terms of the GNU General Public License as published by
-# the Free Software Foundation, either version 3 of the License, or
-# (at your option) any later version.
-#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-# GNU General Public License for more details.
-#
-# You should have received a copy of the GNU General Public License
-# along with this program.  If not, see <http://www.gnu.org/licenses/>
-
-__author__ = "Cedric Bonhomme"
-__version__ = "$Revision: 0.7 $"
-__date__ = "$Date: 2016/03/13 $"
-__revision__ = "$Date: 2019/05/31 $"
-__license__ = "GPLv3"
-
-from typing import IO, Iterator, Union
-
-from PIL import Image
-
-from stegano import tools
-
-
-def hide(
-    input_image: Union[str, IO[bytes]],
-    message: str,
-    generator: Iterator[int],
-    shift: int = 0,
-    encoding: str = "UTF-8",
-    auto_convert_rgb: bool = False,
-):
-    """Hide a message (string) in an image with the
-    LSB (Least Significant Bit) technique.
-    """
-    message_length = len(message)
-    assert message_length != 0, "message length is zero"
-
-    img = tools.open_image(input_image)
-
-    if img.mode not in ["RGB", "RGBA"]:
-        if not auto_convert_rgb:
-            print("The mode of the image is not RGB. Mode is {}".format(img.mode))
-            answer = input("Convert the image to RGB ? [Y / n]\n") or "Y"
-            if answer.lower() == "n":
-                raise Exception("Not a RGB image.")
-        img = img.convert("RGB")
-
-    img_list = list(img.getdata())
-    width, height = img.size
-    index = 0
-
-    message = str(message_length) + ":" + str(message)
-    message_bits = "".join(tools.a2bits_list(message, encoding))
-    message_bits += "0" * ((3 - (len(message_bits) % 3)) % 3)
-
-    npixels = width * height
-    len_message_bits = len(message_bits)
-    if len_message_bits > npixels * 3:
-        raise Exception(
-            "The message you want to hide is too long: {}".format(message_length)
-        )
-    while shift != 0:
-        next(generator)
-        shift -= 1
-
-    while index + 3 <= len_message_bits:
-        generated_number = next(generator)
-        r, g, b, *a = img_list[generated_number]
-
-        # Change the Least Significant Bit of each colour component.
-        r = tools.setlsb(r, message_bits[index])
-        g = tools.setlsb(g, message_bits[index + 1])
-        b = tools.setlsb(b, message_bits[index + 2])
-
-        # Save the new pixel
-        if img.mode == "RGBA":
-            img_list[generated_number] = (r, g, b, *a)
-        else:
-            img_list[generated_number] = (r, g, b)
-
-        index += 3
-
-    # create empty new image of appropriate format
-    encoded = Image.new(img.mode, (img.size))
-
-    # insert saved data into the image
-    encoded.putdata(img_list)
-
-    return encoded
-
-
-def reveal(
-    input_image: Union[str, IO[bytes]],
-    generator: Iterator[int],
-    shift: int = 0,
-    encoding: str = "UTF-8",
-):
-    """Find a message in an image (with the LSB technique)."""
-    img = tools.open_image(input_image)
-    img_list = list(img.getdata())
-    width, height = img.size
-    buff, count = 0, 0
-    bitab = []
-    limit = None
-
-    while shift != 0:
-        next(generator)
-        shift -= 1
-
-    while True:
-        generated_number = next(generator)
-        # color = [r, g, b]
-        for color in img_list[generated_number][:3]:  # ignore the alpha
-            buff += (color & 1) << (tools.ENCODINGS[encoding] - 1 - count)
-            count += 1
-            if count == tools.ENCODINGS[encoding]:
-                bitab.append(chr(buff))
-                buff, count = 0, 0
-                if bitab[-1] == ":" and limit is None:
-                    if "".join(bitab[:-1]).isdigit():
-                        limit = int("".join(bitab[:-1]))
-                    else:
-                        raise IndexError("Impossible to detect message.")
-        if len(bitab) - len(str(limit)) - 1 == limit:
-            return "".join(bitab)[len(str(limit)) + 1 :]
diff --git a/stegano/tools.py b/stegano/tools.py
index 6ddeac3..78fb9f3 100755
--- a/stegano/tools.py
+++ b/stegano/tools.py
@@ -115,3 +115,104 @@ def open_image(fname_or_instance: Union[str, IO[bytes]]):
         return fname_or_instance
 
     return Image.open(fname_or_instance)
+
+
+class Hider:
+    def __init__(
+        self,
+        input_image: Union[str, IO[bytes]],
+        message: str,
+        encoding: str = "UTF-8",
+        auto_convert_rgb: bool = False,
+    ):
+        self._index = 0
+
+        message_length = len(message)
+        assert message_length != 0, "message length is zero"
+
+        image = open_image(input_image)
+
+        if image.mode not in ["RGB", "RGBA"]:
+            if not auto_convert_rgb:
+                print("The mode of the image is not RGB. Mode is {}".format(image.mode))
+                answer = input("Convert the image to RGB ? [Y / n]\n") or "Y"
+                if answer.lower() == "n":
+                    raise Exception("Not a RGB image.")
+
+            image = image.convert("RGB")
+
+        self.encoded_image = image.copy()
+        image.close()
+
+        message = str(message_length) + ":" + str(message)
+        self._message_bits = "".join(a2bits_list(message, encoding))
+        self._message_bits += "0" * ((3 - (len(self._message_bits) % 3)) % 3)
+
+        width, height = self.encoded_image.size
+        npixels = width * height
+        self._len_message_bits = len(self._message_bits)
+
+        if self._len_message_bits > npixels * 3:
+            raise Exception(
+                "The message you want to hide is too long: {}".format(message_length)
+            )
+
+    def encode_another_pixel(self):
+        return True if self._index + 3 <= self._len_message_bits else False
+
+    def encode_pixel(self, coordinate: tuple):
+        # Get the colour component.
+        r, g, b, *a = self.encoded_image.getpixel(coordinate)
+
+        # Change the Least Significant Bit of each colour component.
+        r = setlsb(r, self._message_bits[self._index])
+        g = setlsb(g, self._message_bits[self._index + 1])
+        b = setlsb(b, self._message_bits[self._index + 2])
+
+        # Save the new pixel
+        if self.encoded_image.mode == "RGBA":
+            self.encoded_image.putpixel(coordinate, (r, g, b, *a))
+        else:
+            self.encoded_image.putpixel(coordinate, (r, g, b))
+
+        self._index += 3
+
+
+class Revealer:
+    def __init__(self, encoded_image: Union[str, IO[bytes]], encoding: str = "UTF-8"):
+        self.encoded_image = open_image(encoded_image)
+        self._encoding_length = ENCODINGS[encoding]
+        self._buff, self._count = 0, 0
+        self._bitab = []
+        self._limit = None
+        self.secret_message = ""
+
+    def decode_pixel(self, coordinate: tuple):
+        # pixel = [r, g, b] or [r,g,b,a]
+        pixel = self.encoded_image.getpixel(coordinate)
+
+        if self.encoded_image.mode == "RGBA":
+            pixel = pixel[:3]  # ignore the alpha
+
+        for color in pixel:
+            self._buff += (color & 1) << (self._encoding_length - 1 - self._count)
+            self._count += 1
+
+            if self._count == self._encoding_length:
+                self._bitab.append(chr(self._buff))
+                self._buff, self._count = 0, 0
+
+                if self._bitab[-1] == ":" and self._limit is None:
+                    if "".join(self._bitab[:-1]).isdigit():
+                        self._limit = int("".join(self._bitab[:-1]))
+                    else:
+                        raise IndexError("Impossible to detect message.")
+
+        if len(self._bitab) - len(str(self._limit)) - 1 == self._limit:
+            self.secret_message = "".join(self._bitab)[len(str(self._limit)) + 1 :]
+            self.encoded_image.close()
+
+            return True
+
+        else:
+            return False
diff --git a/tests/test_generators.py b/tests/test_generators.py
index 914291c..a7d7309 100644
--- a/tests/test_generators.py
+++ b/tests/test_generators.py
@@ -30,7 +30,7 @@ import itertools
 import cv2
 import numpy as np
 
-from stegano.lsbset import generators
+from stegano.lsb import generators
 
 
 class TestGenerators(unittest.TestCase):
@@ -151,7 +151,7 @@ class TestGenerators(unittest.TestCase):
         """Test the LFSR generator"""
         with open("./tests/expected-results/LFSR", "r") as f:
             self.assertEqual(
-                tuple(itertools.islice(generators.LFSR(2 ** 8), 256)),
+                tuple(itertools.islice(generators.LFSR(2**8), 256)),
                 tuple(int(line) for line in f),
             )
 
diff --git a/tests/test_lsb.py b/tests/test_lsb.py
index 99ce4e3..a50620c 100644
--- a/tests/test_lsb.py
+++ b/tests/test_lsb.py
@@ -2,7 +2,7 @@
 # -*- coding: utf-8 -*-
 
 # Stegano - Stegano is a pure Python steganography module.
-# Copyright (C) 2010-2022 Cédric Bonhomme - https://www.cedricbonhomme.org
+# Copyright (C) 2010-2022  Cédric Bonhomme - https://www.cedricbonhomme.org
 #
 # For more information : https://git.sr.ht/~cedric/stegano
 #
@@ -20,18 +20,18 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>
 
 __author__ = "Cedric Bonhomme"
-__version__ = "$Revision: 0.3 $"
-__date__ = "$Date: 2016/04/12 $"
-__revision__ = "$Date: 2017/05/04 $"
+__version__ = "$Revision: 0.6 $"
+__date__ = "$Date: 2016/04/13 $"
+__revision__ = "$Date: 2022/01/04 $"
 __license__ = "GPLv3"
 
 import io
 import os
-import base64
 import unittest
 from unittest.mock import patch
 
 from stegano import lsb
+from stegano.lsb import generators
 
 
 class TestLSB(unittest.TestCase):
@@ -40,133 +40,198 @@ class TestLSB(unittest.TestCase):
         Test hiding the empty string.
         """
         with self.assertRaises(AssertionError):
-            lsb.hide("./tests/sample-files/Lenna.png", "")
+            lsb.hide("./tests/sample-files/Lenna.png", "", generators.eratosthenes())
 
-    def test_hide_and_reveal(self):
+    def test_hide_and_reveal_without_generator(self):
         messages_to_hide = ["a", "foo", "Hello World!", ":Python:"]
         for message in messages_to_hide:
             secret = lsb.hide("./tests/sample-files/Lenna.png", message)
             secret.save("./image.png")
 
             clear_message = lsb.reveal("./image.png")
+
+            self.assertEqual(message, clear_message)
+
+    def test_hide_and_reveal_with_eratosthenes(self):
+        messages_to_hide = ["a", "foo", "Hello World!", ":Python:"]
+        for message in messages_to_hide:
+            secret = lsb.hide(
+                "./tests/sample-files/Lenna.png", message, generators.eratosthenes()
+            )
+            secret.save("./image.png")
+
+            clear_message = lsb.reveal("./image.png", generators.eratosthenes())
+
+            self.assertEqual(message, clear_message)
+
+    def test_hide_and_reveal_with_ackermann(self):
+        messages_to_hide = ["foo"]
+        for message in messages_to_hide:
+            secret = lsb.hide(
+                "./tests/sample-files/Lenna.png", message, generators.ackermann(m=3)
+            )
+            secret.save("./image.png")
+
+            clear_message = lsb.reveal("./image.png", generators.ackermann(m=3))
+
+            self.assertEqual(message, clear_message)
+
+    def test_hide_and_reveal_with_ackermann_naive(self):
+        messages_to_hide = ["foo"]
+        for message in messages_to_hide:
+            secret = lsb.hide(
+                "./tests/sample-files/Lenna.png",
+                message,
+                generators.ackermann_naive(m=2),
+            )
+            secret.save("./image.png")
+
+            clear_message = lsb.reveal("./image.png", generators.ackermann_naive(m=2))
+
+            self.assertEqual(message, clear_message)
+
+    def test_hide_and_reveal_with_mersenne(self):
+        messages_to_hide = ["f"]
+        for message in messages_to_hide:
+            secret = lsb.hide(
+                "./tests/sample-files/Lenna.png",
+                message,
+                generators.mersenne(),
+            )
+            secret.save("./image.png")
+
+            clear_message = lsb.reveal("./image.png", generators.mersenne())
+
+            self.assertEqual(message, clear_message)
+
+    def test_hide_and_reveal_with_shi_tomashi(self):
+        messages_to_hide = ["foo bar"]
+        for message in messages_to_hide:
+            secret = lsb.hide(
+                "./tests/sample-files/Lenna.png",
+                message,
+                generators.shi_tomashi("./tests/sample-files/Lenna.png"),
+            )
+            secret.save("./image.png")
+
+            clear_message = lsb.reveal(
+                "./image.png", generators.shi_tomashi("./tests/sample-files/Lenna.png")
+            )
+
+            self.assertEqual(message, clear_message)
+
+    def test_hide_and_reveal_with_shift(self):
+        messages_to_hide = ["a", "foo", "Hello World!", ":Python:"]
+        for message in messages_to_hide:
+            secret = lsb.hide(
+                "./tests/sample-files/Lenna.png", message, generators.eratosthenes(), 4
+            )
+            secret.save("./image.png")
+
+            clear_message = lsb.reveal("./image.png", generators.eratosthenes(), 4)
+
             self.assertEqual(message, clear_message)
 
     def test_hide_and_reveal_UTF32LE(self):
         messages_to_hide = "I love 🍕 and 🍫!"
         secret = lsb.hide(
-            "./tests/sample-files/Lenna.png", messages_to_hide, encoding="UTF-32LE"
+            "./tests/sample-files/Lenna.png",
+            messages_to_hide,
+            generators.eratosthenes(),
+            encoding="UTF-32LE",
         )
         secret.save("./image.png")
 
-        clear_message = lsb.reveal("./image.png", encoding="UTF-32LE")
+        clear_message = lsb.reveal(
+            "./image.png", generators.eratosthenes(), encoding="UTF-32LE"
+        )
         self.assertEqual(messages_to_hide, clear_message)
 
     def test_with_transparent_png(self):
-        messages_to_hide = ["🍕", "a", "foo", "Hello World!", ":Python:"]
+        messages_to_hide = ["a", "foo", "Hello World!", ":Python:"]
         for message in messages_to_hide:
             secret = lsb.hide(
-                "./tests/sample-files/transparent.png", message, encoding="UTF-32LE"
+                "./tests/sample-files/transparent.png",
+                message,
+                generators.eratosthenes(),
             )
             secret.save("./image.png")
 
-            clear_message = lsb.reveal("./image.png", encoding="UTF-32LE")
+            clear_message = lsb.reveal("./image.png", generators.eratosthenes())
 
             self.assertEqual(message, clear_message)
 
     @patch("builtins.input", return_value="y")
     def test_manual_convert_rgb(self, input):
-        message_to_hide = "I love 🍕 and 🍫!"
+        message_to_hide = "Hello World!"
         lsb.hide(
             "./tests/sample-files/Lenna-grayscale.png",
             message_to_hide,
-            encoding="UTF-32LE",
+            generators.eratosthenes(),
         )
 
     @patch("builtins.input", return_value="n")
     def test_refuse_convert_rgb(self, input):
-        message_to_hide = "I love 🍕 and 🍫!"
+        message_to_hide = "Hello World!"
         with self.assertRaises(Exception):
             lsb.hide(
                 "./tests/sample-files/Lenna-grayscale.png",
                 message_to_hide,
-                encoding="UTF-32LE",
+                generators.eratosthenes(),
             )
 
-    def test_auto_convert_rgb(self):
-        message_to_hide = "I love 🍕 and 🍫!"
-        lsb.hide(
-            "./tests/sample-files/Lenna-grayscale.png",
-            message_to_hide,
-            encoding="UTF-32LE",
-            auto_convert_rgb=True,
-        )
-
-    def test_with_text_file(self):
-        text_file_to_hide = "./tests/sample-files/lorem_ipsum.txt"
-        with open(text_file_to_hide) as f:
-            message = f.read()
-        secret = lsb.hide("./tests/sample-files/Lenna.png", message)
-        secret.save("./image.png")
-
-        clear_message = lsb.reveal("./image.png")
-        self.assertEqual(message, clear_message)
-
-    def test_with_binary_file(self):
-        binary_file_to_hide = "./tests/sample-files/free-software-song.ogg"
-        with open(binary_file_to_hide, "rb") as bin_file:
-            encoded_string = base64.b64encode(bin_file.read())
-            message = encoded_string.decode()
-        secret = lsb.hide("./tests/sample-files/Montenach.png", message)
-        secret.save("./image.png")
-
-        clear_message = lsb.reveal("./image.png")
-        clear_message += "==="
-        clear_message = base64.b64decode(clear_message)
-        with open("file1", "wb") as f:
-            f.write(clear_message)
-        with open("file1", "rb") as bin_file:
-            encoded_string = base64.b64encode(bin_file.read())
-            message1 = encoded_string.decode()
-        self.assertEqual(message, message1)
-        try:
-            os.unlink("./file1")
-        except Exception:
-            pass
-
-    def test_with_too_long_message(self):
-        with open("./tests/sample-files/lorem_ipsum.txt") as f:
-            message = f.read()
-        message += message * 2
-        with self.assertRaises(Exception):
-            lsb.hide("./tests/sample-files/Lenna.png", message)
-
-    def test_with_bytes(self):
-        messages_to_hide = ["a", "foo", "Hello World!", ":Python:"]
-
-        for message in messages_to_hide:
-            outputBytes = io.BytesIO()
-            with open("./tests/sample-files/20160505T130442.jpg", "rb") as f:
-                bytes_image = lsb.hide(f, message)
-                bytes_image.save(outputBytes, "PNG")
-                outputBytes.seek(0)
-
-                clear_message = lsb.reveal(outputBytes)
-
-                self.assertEqual(message, clear_message)
-
     def test_with_location_of_image_as_argument(self):
         messages_to_hide = ["Hello World!"]
 
         for message in messages_to_hide:
             outputBytes = io.BytesIO()
-            bytes_image = lsb.hide("./tests/sample-files/20160505T130442.jpg", message)
+            bytes_image = lsb.hide(
+                "./tests/sample-files/20160505T130442.jpg",
+                message,
+                generators.identity(),
+            )
             bytes_image.save(outputBytes, "PNG")
             outputBytes.seek(0)
 
-            clear_message = lsb.reveal(outputBytes)
+            clear_message = lsb.reveal(outputBytes, generators.identity())
 
             self.assertEqual(message, clear_message)
 
+    def test_auto_convert_rgb(self):
+        message_to_hide = "Hello World!"
+        lsb.hide(
+            "./tests/sample-files/Lenna-grayscale.png",
+            message_to_hide,
+            generators.eratosthenes(),
+            auto_convert_rgb=True,
+        )
+
+    def test_with_too_long_message(self):
+        with open("./tests/sample-files/lorem_ipsum.txt") as f:
+            message = f.read()
+        message += message * 2
+        with self.assertRaises(Exception):
+            lsb.hide("./tests/sample-files/Lenna.png", message, generators.identity())
+
+    def test_hide_and_reveal_with_bad_generator(self):
+        message_to_hide = "Hello World!"
+        secret = lsb.hide(
+            "./tests/sample-files/Lenna.png", message_to_hide, generators.eratosthenes()
+        )
+        secret.save("./image.png")
+
+        with self.assertRaises(IndexError):
+            lsb.reveal("./image.png", generators.identity())
+
+    def test_with_unknown_generator(self):
+        message_to_hide = "Hello World!"
+        with self.assertRaises(AttributeError):
+            lsb.hide(
+                "./tests/sample-files/Lenna.png",
+                message_to_hide,
+                generators.unknown_generator(),
+            )
+
     def tearDown(self):
         try:
             os.unlink("./image.png")
diff --git a/tests/test_lsbset.py b/tests/test_lsbset.py
deleted file mode 100644
index 90658f6..0000000
--- a/tests/test_lsbset.py
+++ /dev/null
@@ -1,207 +0,0 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-
-# Stegano - Stegano is a pure Python steganography module.
-# Copyright (C) 2010-2022  Cédric Bonhomme - https://www.cedricbonhomme.org
-#
-# For more information : https://git.sr.ht/~cedric/stegano
-#
-# This program is free software: you can redistribute it and/or modify
-# it under the terms of the GNU General Public License as published by
-# the Free Software Foundation, either version 3 of the License, or
-# (at your option) any later version.
-#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-# GNU General Public License for more details.
-#
-# You should have received a copy of the GNU General Public License
-# along with this program.  If not, see <http://www.gnu.org/licenses/>
-
-__author__ = "Cedric Bonhomme"
-__version__ = "$Revision: 0.6 $"
-__date__ = "$Date: 2016/04/13 $"
-__revision__ = "$Date: 2022/01/04 $"
-__license__ = "GPLv3"
-
-import io
-import os
-import unittest
-from unittest.mock import patch
-
-from stegano import lsbset
-from stegano.lsbset import generators
-
-
-class TestLSBSet(unittest.TestCase):
-    def test_hide_empty_message(self):
-        """
-        Test hiding the empty string.
-        """
-        with self.assertRaises(AssertionError):
-            lsbset.hide("./tests/sample-files/Lenna.png", "", generators.eratosthenes())
-
-    def test_hide_and_reveal(self):
-        messages_to_hide = ["a", "foo", "Hello World!", ":Python:"]
-        for message in messages_to_hide:
-            secret = lsbset.hide(
-                "./tests/sample-files/Lenna.png", message, generators.eratosthenes()
-            )
-            secret.save("./image.png")
-
-            clear_message = lsbset.reveal("./image.png", generators.eratosthenes())
-
-            self.assertEqual(message, clear_message)
-
-    def test_hide_and_reveal_with_ackermann(self):
-        messages_to_hide = ["foo"]
-        for message in messages_to_hide:
-            secret = lsbset.hide(
-                "./tests/sample-files/Lenna.png", message, generators.ackermann(m=3)
-            )
-            secret.save("./image.png")
-
-            clear_message = lsbset.reveal("./image.png", generators.ackermann(m=3))
-
-            self.assertEqual(message, clear_message)
-
-    def test_hide_and_reveal_with_shi_tomashi(self):
-        messages_to_hide = ["foo bar"]
-        for message in messages_to_hide:
-            secret = lsbset.hide(
-                "./tests/sample-files/Lenna.png",
-                message,
-                generators.shi_tomashi("./tests/sample-files/Lenna.png"),
-            )
-            secret.save("./image.png")
-
-            clear_message = lsbset.reveal(
-                "./image.png", generators.shi_tomashi("./tests/sample-files/Lenna.png")
-            )
-
-            self.assertEqual(message, clear_message)
-
-    def test_hide_and_reveal_with_shift(self):
-        messages_to_hide = ["a", "foo", "Hello World!", ":Python:"]
-        for message in messages_to_hide:
-            secret = lsbset.hide(
-                "./tests/sample-files/Lenna.png", message, generators.eratosthenes(), 4
-            )
-            secret.save("./image.png")
-
-            clear_message = lsbset.reveal("./image.png", generators.eratosthenes(), 4)
-
-            self.assertEqual(message, clear_message)
-
-    def test_hide_and_reveal_UTF32LE(self):
-        messages_to_hide = "I love 🍕 and 🍫!"
-        secret = lsbset.hide(
-            "./tests/sample-files/Lenna.png",
-            messages_to_hide,
-            generators.eratosthenes(),
-            encoding="UTF-32LE",
-        )
-        secret.save("./image.png")
-
-        clear_message = lsbset.reveal(
-            "./image.png", generators.eratosthenes(), encoding="UTF-32LE"
-        )
-        self.assertEqual(messages_to_hide, clear_message)
-
-    def test_with_transparent_png(self):
-        messages_to_hide = ["a", "foo", "Hello World!", ":Python:"]
-        for message in messages_to_hide:
-            secret = lsbset.hide(
-                "./tests/sample-files/transparent.png",
-                message,
-                generators.eratosthenes(),
-            )
-            secret.save("./image.png")
-
-            clear_message = lsbset.reveal("./image.png", generators.eratosthenes())
-
-            self.assertEqual(message, clear_message)
-
-    @patch("builtins.input", return_value="y")
-    def test_manual_convert_rgb(self, input):
-        message_to_hide = "Hello World!"
-        lsbset.hide(
-            "./tests/sample-files/Lenna-grayscale.png",
-            message_to_hide,
-            generators.eratosthenes(),
-        )
-
-    @patch("builtins.input", return_value="n")
-    def test_refuse_convert_rgb(self, input):
-        message_to_hide = "Hello World!"
-        with self.assertRaises(Exception):
-            lsbset.hide(
-                "./tests/sample-files/Lenna-grayscale.png",
-                message_to_hide,
-                generators.eratosthenes(),
-            )
-
-    def test_with_location_of_image_as_argument(self):
-        messages_to_hide = ["Hello World!"]
-
-        for message in messages_to_hide:
-            outputBytes = io.BytesIO()
-            bytes_image = lsbset.hide(
-                "./tests/sample-files/20160505T130442.jpg",
-                message,
-                generators.identity(),
-            )
-            bytes_image.save(outputBytes, "PNG")
-            outputBytes.seek(0)
-
-            clear_message = lsbset.reveal(outputBytes, generators.identity())
-
-            self.assertEqual(message, clear_message)
-
-    def test_auto_convert_rgb(self):
-        message_to_hide = "Hello World!"
-        lsbset.hide(
-            "./tests/sample-files/Lenna-grayscale.png",
-            message_to_hide,
-            generators.eratosthenes(),
-            auto_convert_rgb=True,
-        )
-
-    def test_with_too_long_message(self):
-        with open("./tests/sample-files/lorem_ipsum.txt") as f:
-            message = f.read()
-        message += message * 2
-        with self.assertRaises(Exception):
-            lsbset.hide(
-                "./tests/sample-files/Lenna.png", message, generators.identity()
-            )
-
-    def test_hide_and_reveal_with_bad_generator(self):
-        message_to_hide = "Hello World!"
-        secret = lsbset.hide(
-            "./tests/sample-files/Lenna.png", message_to_hide, generators.eratosthenes()
-        )
-        secret.save("./image.png")
-
-        with self.assertRaises(IndexError):
-            lsbset.reveal("./image.png", generators.identity())
-
-    def test_with_unknown_generator(self):
-        message_to_hide = "Hello World!"
-        with self.assertRaises(AttributeError):
-            lsbset.hide(
-                "./tests/sample-files/Lenna.png",
-                message_to_hide,
-                generators.eratosthene(),  # type: ignore
-            )
-
-    def tearDown(self):
-        try:
-            os.unlink("./image.png")
-        except Exception:
-            pass
-
-
-if __name__ == "__main__":
-    unittest.main()