Use Python's Scipy DCT-II to do 2D or ND DCT

Question

I would like to use scipy's DCT-II since it is already coded and fast. Looking at the doc, it seems it is the 1D implementation. Is it possible to use it in such a way to use it as a 3D implementation? I'm not sure about the math. Are the 2D and 3D implementations the equivalent of multiplying 2 or 3 times the 1D using different dimensions in the calculation?

Answer 1

Basically, the following does the trick:

import numpy as np
from scipy.fftpack import dct, idct

# Lets create a 3D array and fill it with some values
a = np.random.rand(3,3,3)

b = dct(dct(dct(a).transpose(0,2,1)).transpose(1,2,0)).transpose(1,2,0).transpose(0,2,1)

Answer 2

Here's a more general function:

import numpy as np
from scipy.fftpack import dct, idct

def dctn(x, norm="ortho"):
    for i in range(x.ndim):
        x = dct(x, axis=i, norm=norm)
    return x

def idctn(x, norm="ortho"):
    for i in range(x.ndim):
        x = idct(x, axis=i, norm=norm)
    return x

Then:

>>> x = np.random.rand(2, 2, 2)

>>> x
array([[[0.316, 0.927],
        [0.197, 0.936]],

       [[0.832, 0.982],
        [0.768, 0.564]]])

>>> dctn(x)
array([[[ 1.952, -0.459],
        [ 0.209, -0.08 ]],

       [[-0.272, -0.496],
        [-0.132,  0.171]]])

>>> np.all(np.isclose(x, idctn(dctn(x))))
True

---

To make the output the exact same as @macrocosme's, set norm=None and run this at the end of dctn:

np.moveaxis(x, range(x.ndim), (-1, range(x.ndim - 1)))

---

As a further example, we can verify that it works the same as JPEG's 8x8 DCT-II with Wikipedia's calculations:

>>> x = np.array([
...     [52, 55, 61, 66,  70,  61,  64, 73],
...     [63, 59, 55, 90,  109, 85,  69, 72],
...     [62, 59, 68, 113, 144, 104, 66, 73],
...     [63, 58, 71, 122, 154, 106, 70, 69],
...     [67, 61, 68, 104, 126, 88,  68, 70],
...     [79, 65, 60, 70,  77,  68,  58, 75],
...     [85, 71, 64, 59,  55,  61,  65, 83],
...     [87, 79, 69, 68,  65,  76,  78, 94],
... ]) - 128

>>> np.set_printoptions(precision=2, floatmode="fixed", suppress=True)

>>> print(dctn(x))
[[-415.38  -30.19  -61.20   27.24   56.12  -20.10   -2.39    0.46]
 [   4.47  -21.86  -60.76   10.25   13.15   -7.09   -8.54    4.88]
 [ -46.83    7.37   77.13  -24.56  -28.91    9.93    5.42   -5.65]
 [ -48.53   12.07   34.10  -14.76  -10.24    6.30    1.83    1.95]
 [  12.12   -6.55  -13.20   -3.95   -1.88    1.75   -2.79    3.14]
 [  -7.73    2.91    2.38   -5.94   -2.38    0.94    4.30    1.85]
 [  -1.03    0.18    0.42   -2.42   -0.88   -3.02    4.12   -0.66]
 [  -0.17    0.14   -1.07   -4.19   -1.17   -0.10    0.50    1.68]]