jpeg reader and generator

data.py

+import os
+import jpegio as jio
+import numpy as np
+
+from utils import decompress_structure
+
+
 def get_train_test_generator(dir_path, train_size, stego_percentage):
     """
     Return two generators. One for the train cover images, the second for the test cover images and a third for the
@@ -7,7 +14,25 @@ def get_train_test_generator(dir_path, train_size, stego_percentage):
     :param stego_percentage: percentage of the stego images among the test images
     :return: three generators.
     """
-    pass
+
+    # TODO: identify why some images raise a "Premature end of JPEG file" and correct it.
+    # TODO: Read .pgm and compress them in .jpeg.
+
+    def img_generator(dir_path, names):
+        for name in names:
+            path = os.path.join(dir_path, name)
+            tmp = jio.read(path)
+            img = decompress_structure(tmp)[:, :, 0].astype(np.float32)
+            yield img
+
+    names = os.listdir(dir_path)
+    n = len(names)
+    n_train = int(n * train_size)
+    n_normal = int((n - n_train) * (1 - stego_percentage))
+
+    return img_generator(dir_path, names[: n_train]), \
+           img_generator(dir_path, names[n_train: n_train + n_normal]), \
+           img_generator(dir_path, names[n_train + n_normal:])
 
 
 def embed_images(img_generator, payload):

utils.py

+import numpy as np
+from numpy.lib.stride_tricks import as_strided
+from scipy import fftpack
+
+def block_view(A, block=(8, 8)):
+    """Provide a 2D block view to 2D array. No error checking made.
+    Therefore, meaningful (as implemented) only for blocks strictly
+    compatible with the shape of A."""
+    # simple shape and strides computations may seem at first strange
+    # unless one is able to recognize the 'tuple additions' involved ;-)
+    shape = (A.shape[0] // block[0], A.shape[1] // block[1]) + block
+    strides = (block[0] * A.strides[0], block[1] * A.strides[1]) + A.strides
+    return as_strided(A, shape=shape, strides=strides)
+
+
+def segmented_stride(M, fun, blk_size=(8, 8), overlap=(0, 0)):
+    # This is some complex function of blk_size and M.shape
+    B = block_view(M, block=blk_size)
+    B[:, :, :, :] = fun(B)
+    return M
+
+
+def decompress_structure(S, grayscale=True):
+    # Decompress DCT coefficients C using quantization table Q
+    H = S.coef_arrays[0].shape[0]
+    W = S.coef_arrays[0].shape[1]
+    if grayscale:
+        n = 1
+    else:
+        n = len(S.coef_arrays)
+    assert H % 8 == 0, 'Wrong image size'
+    assert W % 8 == 0, 'Wrong image size'
+    I = np.zeros((H, W, n), dtype=np.float64)  # Returns Y, Cb and Cr
+    for i in range(n):
+        Q = S.quant_tables[S.comp_info[i].quant_tbl_no]
+        # this multiplication is done on integers
+        fun = lambda x: np.multiply(x, Q)
+        C = np.float64(segmented_stride(S.coef_arrays[i], fun))
+        fun = lambda x: fftpack.idct(fftpack.idct(x, norm='ortho', axis=2), norm='ortho', axis=3) + 128
+        I[:, :, i] = segmented_stride(C, fun)
+    return I
\ No newline at end of file