Updated image Steg to allow for multiple chosen bits as bits to use for hiding

Image_Steg/img_steg.py

@@ -3,84 +3,26 @@ import numpy as np
 
 
 class img_steg:
-    def __init__(self, image_name: str = "image.png", bit_to_hide: int = 2):
+    def __init__(self, image_name: str = "image.png", bit_to_hide: list[int] = None):
         """
         Initialize the class
         :param image_name:  Name of the image to encode or decode
         :type image_name: str
-        :param bit_to_hide: Bit to hide the data in (1 - LSB to 8 - MSB)
+        :param bit_to_hide: List of bit position to hide the data (LSB is 1, MSB is 8)
-        :type bit_to_hide: int
+        :type bit_to_hide: list[int]
         """
         self.image_name = image_name
-        self.bit_to_hide = 8 - bit_to_hide
+        self.bit_to_hide = [8 - bit_pos for bit_pos in bit_to_hide]
-        self.delimiter = "\n"
+        self.delimiter = "====="
-
-    def to_bin(self, data: str) -> str | list[str]:
-        """
-        Convert data to binary format as string
-        Args:
-            data: String
-
-        Returns:
-            Binary Data: String | List[String]
-        """
-        if isinstance(data, str):
-            return ''.join([format(ord(i), "08b") for i in data])
-        elif isinstance(data, bytes) or isinstance(data, np.ndarray):
-            return [format(i, "08b") for i in data]
-        elif isinstance(data, int) or isinstance(data, np.unit8):
-            return format(data, "08b")
-        else:
-            raise TypeError("Type not supported")
-
-    def from_bin(self, data: str) -> str:
-        """
-        Convert binary `data` back to the original format
-        Args:
-            data: String
-
-        Returns:
-            Original Data: String
-        """
-        return ''.join([chr(int(data[i:i + 8], 2)) for i in range(0, len(data), 8)])
-
-    def decode(self) -> str:
-        """
-        Decode the data hidden in the image
-        Args:
-            None
-        :return: Decoded Data: String
-        """
-        print("[+] Decoding...")
-        image = cv2.imread(self.image_name)  # read image
-        binary_data = ""
-        bit_to_hide = self.bit_to_hide
-        for row in image:
-            for pixel in row:
-                r, g, b = self.to_bin(pixel)
-                binary_data += r[bit_to_hide]
-                binary_data += g[bit_to_hide]
-                binary_data += b[bit_to_hide]
-        # Split by 8 bits
-        all_bytes = [binary_data[i: i + 8] for i in range(0, len(binary_data), 8)]
-        # Convert from bits to characters
-        decoded_data = ""
-        for byte in all_bytes:
-            decoded_data += chr(int(byte, 2))
-            if decoded_data[-len(self.delimiter):] == self.delimiter:
-                break
-
-        return decoded_data[:-len(self.delimiter)]
 
     def encode(self, secret_data: str = "Hello World") -> np.ndarray:
         """
-       Encode the secret data in the image
+        Encode the secret data in the image
-       Args:
+        :param secret_data: Data to hide in the image
-           secret_data: str
+        :type secret_data: str
-
+        :return: Encoded Image
-       Returns:
+        :rtype: np.ndarray
-           Encoded Image: np.ndarray
+        """
-       """
         image = cv2.imread(self.image_name)  # read image
         n_bytes = image.shape[0] * image.shape[1] * 3 // 8  # Max bytes to encode
         print("[*] Maximum bytes to encode:", n_bytes)
@@ -97,43 +39,119 @@ class img_steg:
         for row in image:
             for pixel in row:
                 r, g, b = self.to_bin(pixel)  # Convert RGB Values to binary format
-                if data_index < data_len:
+                for bit_pos in bit_to_hide:
-                    r = list(r)
+                    if data_index < data_len:
-                    r[bit_to_hide] = binary_secret_data[data_index]  # hide data into least significant bit of red pixel
+                        r = list(r)
-                    pixel[0] = int(''.join(r), 2)
+                        r[bit_pos] = binary_secret_data[
-                    data_index += 1
+                            data_index]  # hide data into specified bit position of red pixel
-                if data_index < data_len:
+                        pixel[0] = int(''.join(r), 2)
-                    g = list(g)
+                        data_index += 1
-                    g[bit_to_hide] = binary_secret_data[data_index]
+                    if data_index < data_len:
-                    pixel[1] = int(''.join(g), 2)
+                        g = list(g)
-                    data_index += 1
+                        g[bit_pos] = binary_secret_data[
-                if data_index < data_len:
+                            data_index]
-                    b = list(b)
+                        pixel[1] = int(''.join(g), 2)
-                    b[bit_to_hide] = binary_secret_data[data_index]
+                        data_index += 1
-                    pixel[2] = int(''.join(b), 2)
+                    if data_index < data_len:
-                    data_index += 1
+                        b = list(b)
+                        b[bit_pos] = binary_secret_data[
+                            data_index]
+                        pixel[2] = int(''.join(b), 2)
+                        data_index += 1
+                    if data_index >= data_len:
+                        break
+
+                pixel[0] = int(''.join(r), 2)  # convert modified binary back to integer for red pixel
+                pixel[1] = int(''.join(g), 2)  # convert modified binary back to integer for green pixel
+                pixel[2] = int(''.join(b), 2)  # convert modified binary back to integer for blue pixel
+
                 if data_index >= data_len:
                     break
         return image
 
+    def decode(self) -> str:
+        """
+        Decode the data hidden in the image
+
+        :return: Decoded Data: String
+        """
+        print("[+] Decoding...")
+        image = cv2.imread(self.image_name)  # read image
+        binary_data = ""
+        bit_to_hide = self.bit_to_hide
+        for row in image:
+            for pixel in row:
+                r, g, b = self.to_bin(pixel)
+                for bit_pos in bit_to_hide:
+                    binary_data += r[bit_pos]  # retrieve data from specified bit position of red pixel
+                    binary_data += g[bit_pos]  # retrieve data from specified bit position of green pixel
+                    binary_data += b[bit_pos]  # retrieve data from specified bit position of blue pixel
+
+        # Split by 8 bits
+        all_bytes = [binary_data[i: i + 8] for i in range(0, len(binary_data), 8)]
+        # Convert from bits to characters
+        decoded_data = ""
+        for byte in all_bytes:
+            decoded_data += chr(int(byte, 2))
+            if decoded_data[-len(self.delimiter):] == self.delimiter:
+                break
+
+        return decoded_data[:-len(self.delimiter)]
+
+    def to_bin(self, data: str) -> str | list[str]:
+        """
+        Convert data to binary format as string
+
+        :param data: Data to convert
+        :type data: str
+        :return: Binary Data
+        :rtype: str | list[str]
+        """
+        if isinstance(data, str):
+            return ''.join([format(ord(i), "08b") for i in data])
+        elif isinstance(data, bytes) or isinstance(data, np.ndarray):
+            return [format(i, "08b") for i in data]
+        elif isinstance(data, int) or isinstance(data, np.unit8):
+            return format(data, "08b")
+        else:
+            raise TypeError("Type not supported")
+
+    def from_bin(self, data: str) -> str:
+        """
+        # UNUSED
+        Convert binary `data` back to the original format
+
+        :param data: Binary Data
+        :type data: str
+        :return: Original Data
+        :rtype: str
+        """
+        return ''.join([chr(int(data[i:i + 8], 2)) for i in range(0, len(data), 8)])
+
 
 def main():
+    # Variables
+    image_name = "pokemon.png"
+    encoded_image_name = "encoded_image.png"
+    secret_data = ""
+    with open("../Txt_Steg/test.txt", "r") as f:
+        secret_data = f.read()
+    bit_to_hide = [1, 2, 3, 4, 5, 6]
+
     print("Welcome to Image Steganography")
     print("1. Encode\n2. Decode\n3. Exit")
     choice = int(input("Enter your choice: "))
+
     if choice == 1:
-        image_name = input("Enter name of image to encode: ")
+        encoded_image = img_steg(image_name=image_name, bit_to_hide=bit_to_hide).encode(secret_data)
-        secret_data = input("Enter data to encode: ")
-        bit_to_hide = int(input("Enter bit to hide (1 - LSB to 8 - MSB): "))
-        encoded_image = img_steg(image_name, bit_to_hide).encode(secret_data)
         extension = image_name.split(".")[-1]
         cv2.imwrite("encoded_image." + extension, encoded_image)
         print("Image Encoded Successfully")
     elif choice == 2:
-        image_name = input("Enter name of image to decode: ")
+        decoded_data = img_steg(image_name=image_name, bit_to_hide=bit_to_hide).decode()
-        bit_to_hide = int(input("Enter bit to hide (1 - LSB to 8 - MSB): "))
-        decoded_data = img_steg(image_name, bit_to_hide).decode()
         print("Decoded Data:", decoded_data)
+    else:
+        exit(0)
 
 
 if __name__ == "__main__":
@@ -151,7 +169,7 @@ if __name__ == "__main__":
 # Initialise class with image name and bit to hide
 #
 # image_name = image name with extension (String)
-# bit_to_hide = bit to hide the data in (1 - LSB to 8 - MSB) (int)
+# bit_to_hide = [1, 2, 3] (Please Sort the list in ascending order)
 #
 # To encode:
 # encoded_image = img_steg(image_name, bit_to_hide).encode(secret_data)