{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Faster YC<sub>b</sub>C<sub>r</sub> Decoding\n\nAs discussed in the forums, most methods of loading a JPEG image output an RGB array. Converting this back to YC<sub>b</sub>C<sub>r</sub> can result in tiny errors which could be significant since the stenographic techniques used here rely on tiny changes in DCT space.\n\nRémi Cogranne has kindly provided us a method to decode from JPEG directly to YCbCr using `jpegio`, skipping the RGB step. https://www.kaggle.com/remicogranne/jpeg-explanations\n\nThis method can be a bit slow as discussed [here](https://www.kaggle.com/c/alaska2-image-steganalysis/discussion/150359#845590).\n\nLet's use vectorised operations in NumPy to speed this up","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install git+https://github.com/dwgoon/jpegio\nimport jpegio as jio","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import numpy as np\nfrom pathlib import Path\nfrom PIL import Image\nimport matplotlib.pyplot as plt\n\nINPUT_PATH = Path('/kaggle/input/alaska2-image-steganalysis/')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"files_cover = [x for x in (INPUT_PATH / 'Cover').glob('**/*') if x.is_file()]\nfiles_cover.sort()\n\nlen(files_cover)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# The original version","execution_count":null},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"def JPEGdecompressYCbCr(path):\n    jpegStruct = jio.read(str(path))\n\n    nb_colors = len(jpegStruct.coef_arrays)\n\n    [Col, Row] = np.meshgrid(range(8), range(8))\n    T = 0.5 * np.cos(np.pi * (2 * Col + 1) * Row / (2 * 8))\n    T[0, :] = T[0, :] / np.sqrt(2)\n\n    sz = np.array(jpegStruct.coef_arrays[0].shape)\n\n    imDecompressYCbCr = np.zeros([sz[0], sz[1], nb_colors])\n    szDct = (sz / 8).astype(\"int\")\n\n    for ColorChannel in range(nb_colors):\n        tmpPixels = np.zeros(sz)\n\n        DCTcoefs = jpegStruct.coef_arrays[ColorChannel]\n        if ColorChannel == 0:\n            QM = jpegStruct.quant_tables[ColorChannel]\n        else:\n            QM = jpegStruct.quant_tables[1]\n\n        for idxRow in range(szDct[0]):\n            for idxCol in range(szDct[1]):\n                D = DCTcoefs[\n                    idxRow * 8 : (idxRow + 1) * 8, idxCol * 8 : (idxCol + 1) * 8\n                ]\n                tmpPixels[\n                    idxRow * 8 : (idxRow + 1) * 8, idxCol * 8 : (idxCol + 1) * 8\n                ] = np.dot(np.transpose(T), np.dot(QM * D, T))\n        imDecompressYCbCr[:, :, ColorChannel] = tmpPixels\n        \n    return imDecompressYCbCr.astype(np.float32)","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"img_v1 = JPEGdecompressYCbCr(files_cover[0])\nimg_v1.shape","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Vectorised version\nThe nested `for` loops are causing a large slowdown. We can try two methods of speeding this up:\n* Using `np.einsum`\n* Using matrix multiplication","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def JPEGdecompressYCbCr_v2(path):\n    jpegStruct = jio.read(str(path))\n\n    [col, row] = np.meshgrid(range(8), range(8))\n    T = 0.5 * np.cos(np.pi * (2 * col + 1) * row / (2 * 8))\n    T[0, :] = T[0, :] / np.sqrt(2)\n\n    img_dims = np.array(jpegStruct.coef_arrays[0].shape)\n    n_blocks = img_dims // 8\n    broadcast_dims = (n_blocks[0], 8, n_blocks[1], 8)\n    \n    YCbCr = []\n    for i, dct_coeffs, in enumerate(jpegStruct.coef_arrays):\n\n        if i == 0:\n            QM = jpegStruct.quant_tables[i]\n        else:\n            QM = jpegStruct.quant_tables[1]\n        \n        t = np.broadcast_to(T.reshape(1, 8, 1, 8), broadcast_dims)\n        qm = np.broadcast_to(QM.reshape(1, 8, 1, 8), broadcast_dims)\n        dct_coeffs = dct_coeffs.reshape(broadcast_dims)\n        \n        a = np.einsum('ijkl,ilkm->ijkm', qm * dct_coeffs, t)\n        b = np.einsum('ijkl,ilkm->ijkm', np.transpose(t, axes=(0, 3, 2, 1)), a)\n        YCbCr.append(b.reshape(img_dims))\n                    \n    return np.stack(YCbCr, -1).astype(np.float32)\n\n\ndef JPEGdecompressYCbCr_v3(path):\n    jpegStruct = jio.read(str(path))\n\n    [col, row] = np.meshgrid(range(8), range(8))\n    T = 0.5 * np.cos(np.pi * (2 * col + 1) * row / (2 * 8))\n    T[0, :] = T[0, :] / np.sqrt(2)\n\n    img_dims = np.array(jpegStruct.coef_arrays[0].shape)\n    n_blocks = img_dims // 8\n    broadcast_dims = (n_blocks[0], 8, n_blocks[1], 8)\n    \n    YCbCr = []\n    for i, dct_coeffs, in enumerate(jpegStruct.coef_arrays):\n\n        if i == 0:\n            QM = jpegStruct.quant_tables[i]\n        else:\n            QM = jpegStruct.quant_tables[1]\n        \n        t = np.broadcast_to(T.reshape(1, 8, 1, 8), broadcast_dims)\n        qm = np.broadcast_to(QM.reshape(1, 8, 1, 8), broadcast_dims)\n        dct_coeffs = dct_coeffs.reshape(broadcast_dims)\n        \n        a = np.transpose(t, axes=(0, 2, 3, 1))\n        b = (qm * dct_coeffs).transpose(0,2,1,3)\n        c = t.transpose(0,2,1,3)\n                \n        z = a @ b @ c\n        z = z.transpose(0,2,1,3)\n        YCbCr.append(z.reshape(img_dims))\n                    \n    return np.stack(YCbCr, -1).astype(np.float32)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"img_v2 = JPEGdecompressYCbCr_v2(files_cover[0])  # np.einsum version\nimg_v3 = JPEGdecompressYCbCr_v3(files_cover[0])  # matrix multiplication version\nimg_v2.shape, img_v3.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Check if they are the same\nnp.allclose(img_v1, img_v2), np.allclose(img_v1, img_v3)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Speed comparison","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"%%timeit -n 100\nJPEGdecompressYCbCr(files_cover[0])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%timeit -n 100\nJPEGdecompressYCbCr_v2(files_cover[0])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%timeit -n 100\nJPEGdecompressYCbCr_v3(files_cover[0])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Nice! Quite a decent speedup. Looks like `JPEGdecompressYCbCr_v3` which uses raw matrix multiplication is faster than the Einstein summation method\n\nLet's compare the speed to PIL just for fun","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"%%timeit -n 100\nImage.open(files_cover[0]).convert('YCbCr')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"markdown","source":"# Comparing to the PIL conversion\nNote that to directly compare the PIL array we need to convert fron `uint8` to either a signed `int` or `float` and subtract 128 so that it is zero-centred","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"img_pil = np.array(Image.open(files_cover[0]).convert('YCbCr'), dtype=np.int32) - 128\nimg_pil[0]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"img_v3[0]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"difference = img_pil - img_v3\nprint(f\"Mean difference: {np.mean(difference):0.5f}, Std Dev: {np.std(difference):0.5f}\")\nplt.hist(difference.flatten(), bins=50);","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Not a huge difference, but +/-1 is sometimes all that is needed to hide bits in stenography","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}