Fixed and rename image steg delimeter, able to take in large text data.

img_steg.py

@@ -4,8 +4,16 @@ import numpy as np
 
 class img_steg:
     def __init__(self, image_name: str = "image.png", bit_to_hide: int = 2):
+        """
+        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)
+        :type bit_to_hide: int
+        """
         self.image_name = image_name
-        self.bit_to_hide = bit_to_hide
+        self.bit_to_hide = 8 - bit_to_hide
+        self.delimiter = "\n"
 
     def to_bin(self, data: str) -> str | list[str]:
         """
@@ -36,16 +44,17 @@ class img_steg:
         """
         return ''.join([chr(int(data[i:i + 8], 2)) for i in range(0, len(data), 8)])
 
-    def decode(self):
+    def decode(self) -> str:
         """
         Decode the data hidden in the image
-        Returns:
-            Decoded Data: str
+        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 - 1
+        bit_to_hide = self.bit_to_hide
         for row in image:
             for pixel in row:
                 r, g, b = self.to_bin(pixel)
@@ -56,8 +65,10 @@ class img_steg:
         all_bytes = [binary_data[i: i + 8] for i in range(0, len(binary_data), 8)]
         # Convert from bits to characters
         decoded_data = self.from_bin(''.join(all_bytes))
-        if decoded_data.endswith("===="):
+        if decoded_data.endswith(self.delimiter):
             decoded_data = decoded_data[:-4]
+
+        decoded_data = decoded_data.split(self.delimiter)[0]
         return decoded_data
 
     def encode(self, secret_data: str = "Hello World") -> np.ndarray:
@@ -72,14 +83,13 @@ class img_steg:
         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)
-        secret_data += "===="  # Stopping Criteria
+        secret_data += self.delimiter  # Add delimiter at the end of data
         if len(secret_data) > n_bytes:
             raise ValueError("[!] Insufficient bytes, need a bigger image or less data")
         print("[*] Encoding Data...")
 
         # Convert bit to hide position
-        # bit_to_hide = 2 - self.bit_to_hide
-        bit_to_hide = self.bit_to_hide - 1
+        bit_to_hide = self.bit_to_hide
         data_index = 0
         binary_secret_data = self.to_bin(secret_data)  # Convert data to binary
         data_len = len(binary_secret_data)  # size of data to hide
@@ -113,27 +123,38 @@ def main():
     if choice == 1:
         image_name = input("Enter name of image to encode: ")
         secret_data = input("Enter data to encode: ")
-        bit_to_hide = int(input("Enter bit to hide (1-8): "))
+        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)
-        cv2.imwrite("encoded_image.png", encoded_image)
+        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: ")
-        bit_to_hide = int(input("Enter bit to hide (1-8): "))
+        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) # Still returns special chars and ==== at the end
+        print("Decoded Data:", decoded_data)
 
 
 if __name__ == "__main__":
     main()
-    
-# FOR REFERENCE
-# Welcome to Image Steganography
-# 1. Encode
-# 2. Decode
-# 3. Exit
-# Enter your choice: 2
-# Enter name of image to decode: encoded_image.png
-# Enter bit to hide (1-8): 5
-# [+] Decoding...
-# Decoded Data: abc123abcasdf====
+
+##############################
+#  _   _
+# | | | |___  __ _  __ _  ___
+# | | | / __|/ _` |/ _` |/ _ \
+# | |_| \__ \ (_| | (_| |  __/
+#  \___/|___/\__,_|\__, |\___|
+#                  |___/
+##############################
+#
+# 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)
+#
+# To encode:
+# encoded_image = img_steg(image_name, bit_to_hide).encode(secret_data)
+#
+# To decode:
+# decoded_data = img_steg(image_name, bit_to_hide).decode()
+# Both are string types