Addition of optional parameter save_stages

1) An optional parameter named save_stages is added to the constructor of the ArithmeticEncoding class. If True, then the intervals of each stage are saved in a list. Note that setting save_stages=True may cause memory overflow if the message is large
2) The decoded message returned by the decode() method is always a list.
3) The order of the returned values by the the encode() and decode() method changed.
4) Added an example in the example_image.py script to encode and decode an image.

Author: ahmedfgad

example.py

@@ -1,27 +1,25 @@
 import pyae
-# from decimal import getcontext
 
 # Example for encoding a simple text message using the PyAE module.
 
 frequency_table = {"a": 2,
                    "b": 7,
                    "c": 1}
 
-AE = pyae.ArithmeticEncoding(frequency_table)
-
-# Default precision is 28. Change it to do arithmetic operations with larger/smaller numbers.
-# getcontext().prec = 28
+AE = pyae.ArithmeticEncoding(frequency_table=frequency_table,
+                            save_stages=True)
 
 original_msg = "abc"
 print("Original Message: {msg}".format(msg=original_msg))
 
-encoder, encoded_msg = AE.encode(msg=original_msg, 
+encoded_msg, encoder = AE.encode(msg=original_msg, 
                                  probability_table=AE.probability_table)
 print("Encoded Message: {msg}".format(msg=encoded_msg))
 
-decoder, decoded_msg = AE.decode(encoded_msg=encoded_msg, 
+decoded_msg, decoder = AE.decode(encoded_msg=encoded_msg, 
                                  msg_length=len(original_msg),
                                  probability_table=AE.probability_table)
 print("Decoded Message: {msg}".format(msg=decoded_msg))
 
-print("Message Decoded Successfully? {result}".format(result=original_msg == decoded_msg))
+decoded_msg = "".join(decoded_msg)
+print("Message Decoded Successfully? {result}".format(result=original_msg == decoded_msg))

example2.py

@@ -0,0 +1,34 @@
+import pyae
+# from decimal import getcontext
+
+# Example for encoding a simple text message using the PyAE module.
+
+# Create the frequency table.
+frequency_table = {"a": 2,
+                   "b": 3,
+                   "c": 1,
+                   "d": 4}
+
+# Create an instance of the ArithmeticEncoding class.
+AE = pyae.ArithmeticEncoding(frequency_table, 
+                             save_stages=True)
+
+# Default precision is 28. Change it to do arithmetic operations with larger/smaller numbers.
+# getcontext().prec = 28
+
+original_msg = "bdab"
+print("Original Message: {msg}".format(msg=original_msg))
+
+# Encode the message
+encoded_msg, encoder = AE.encode(msg=original_msg, 
+                                 probability_table=AE.probability_table)
+print("Encoded Message: {msg}".format(msg=encoded_msg))
+
+# Decode the message
+decoded_msg, decoder = AE.decode(encoded_msg=encoded_msg, 
+                                 msg_length=len(original_msg),
+                                 probability_table=AE.probability_table)
+print("Decoded Message: {msg}".format(msg=decoded_msg))
+
+decoded_msg = "".join(decoded_msg)
+print("Message Decoded Successfully? {result}".format(result=original_msg == decoded_msg))

example_image.py

@@ -0,0 +1,48 @@
+import scipy.io
+import pyae
+import numpy
+import matplotlib.pyplot
+
+# Change the precision to a bigger value
+from decimal import getcontext
+getcontext().prec = 10000
+
+# Read an image.
+im = scipy.misc.face(gray=True)
+
+# Just work on a small part to save time. The larger the image, the more time consumed.
+im = im[:50, :50]
+
+# Convert the image into a 1D vector.
+msg = im.flatten()
+
+# Create the frequency table based on its hitogram.
+hist, bin_edges = numpy.histogram(a=im,
+                                  bins=range(0, 257))
+frequency_table = {key: value for key, value in zip(bin_edges[0:256], hist)}
+
+# Create an instance of the ArithmeticEncoding class.
+AE = pyae.ArithmeticEncoding(frequency_table=frequency_table)
+
+# Encode the message
+encoded_msg, _ = AE.encode(msg=msg, 
+                        probability_table=AE.probability_table)
+
+# Decode the message
+decoded_msg, _ = AE.decode(encoded_msg=encoded_msg, 
+                        msg_length=len(msg),
+                        probability_table=AE.probability_table)
+
+# Reshape the image to its original shape.
+decoded_msg = numpy.reshape(decoded_msg, im.shape)
+
+# Show the original and decoded images.
+fig, ax = matplotlib.pyplot.subplots(1, 2)
+ax[0].imshow(im, cmap="gray")
+ax[0].set_title("Original Image")
+ax[0].set_xticks([])
+ax[0].set_yticks([])
+ax[1].imshow(decoded_msg, cmap="gray")
+ax[1].set_title("Reconstructed Image")
+ax[1].set_xticks([])
+ax[1].set_yticks([])

pyae.py

@@ -5,7 +5,16 @@ class ArithmeticEncoding:
     ArithmeticEncoding is a class for building the arithmetic encoding.
     """
 
-    def __init__(self, frequency_table):
+    def __init__(self, frequency_table, save_stages=False):
+        """
+        frequency_table: Frequency table as a dictionary where key is the symbol and value is the frequency.
+        save_stages: If True, then the intervals of each stage are saved in a list. Note that setting save_stages=True may cause memory overflow if the message is large
+        """
+        
+        self.save_stages = save_stages
+        if(save_stages == True):
+            print("WARNING: Setting save_stages=True may cause memory overflow if the message is large.")
+
         self.probability_table = self.get_probability_table(frequency_table)
 
     def get_probability_table(self, frequency_table):
@@ -20,13 +29,13 @@ def get_probability_table(self, frequency_table):
 
         return probability_table
 
-    def get_encoded_value(self, encoder):
+    def get_encoded_value(self, last_stage_probs):
         """
         After encoding the entire message, this method returns the single value that represents the entire message.
         """
-        last_stage = list(encoder[-1].values())
+        last_stage_probs = list(last_stage_probs.values())
         last_stage_values = []
-        for sublist in last_stage:
+        for sublist in last_stage_probs:
             for element in sublist:
                 last_stage_values.append(element)
 
@@ -53,9 +62,12 @@ def encode(self, msg, probability_table):
         """
         Encodes a message.
         """
+        
+        # Make sure 
+        msg = list(msg)
 
         encoder = []
-    
+
         stage_min = Decimal(0.0)
         stage_max = Decimal(1.0)
 
@@ -66,22 +78,26 @@ def encode(self, msg, probability_table):
             stage_min = stage_probs[msg_term][0]
             stage_max = stage_probs[msg_term][1]
 
-            encoder.append(stage_probs)
+            if self.save_stages:
+                encoder.append(stage_probs)
 
-        stage_probs = self.process_stage(probability_table, stage_min, stage_max)
-        encoder.append(stage_probs)
+        last_stage_probs = self.process_stage(probability_table, stage_min, stage_max)
+        
+        if self.save_stages:
+            encoder.append(last_stage_probs)
 
-        encoded_msg = self.get_encoded_value(encoder)
+        encoded_msg = self.get_encoded_value(last_stage_probs)
 
-        return encoder, encoded_msg
+        return encoded_msg, encoder
 
     def decode(self, encoded_msg, msg_length, probability_table):
         """
         Decodes a message.
         """
 
         decoder = []
-        decoded_msg = ""
+
+        decoded_msg = []
 
         stage_min = Decimal(0.0)
         stage_max = Decimal(1.0)
@@ -93,13 +109,16 @@ def decode(self, encoded_msg, msg_length, probability_table):
                 if encoded_msg >= value[0] and encoded_msg <= value[1]:
                     break
 
-            decoded_msg = decoded_msg + msg_term
+            decoded_msg.append(msg_term)
+
             stage_min = stage_probs[msg_term][0]
             stage_max = stage_probs[msg_term][1]
 
-            decoder.append(stage_probs)
+            if self.save_stages:
+                decoder.append(stage_probs)
 
-        stage_probs = self.process_stage(probability_table, stage_min, stage_max)
-        decoder.append(stage_probs)
+        if self.save_stages:
+            last_stage_probs = self.process_stage(probability_table, stage_min, stage_max)
+            decoder.append(last_stage_probs)
 
-        return decoder, decoded_msg
+        return decoded_msg, decoder