Two Dimensional Discrete Cosine Transform (DCT-II) in Python
Here’s my first shot at implementing the 2D discrete cosine transform in Python. It seems to work, but let me know in the comments if it’s wrong, or if you know of any easy ways to speed it up.
import numpy
#shortcuts
from numpy import pi
from math import cos,sqrt
def two_dim_DCT(X):
"""2D Discrete Cosine Transform
X should be square 2 dimensional array
Trying to follow:
http://en.wikipedia.org/wiki/Discrete_cosine_transform#Multidimensional_DCTs
http://en.wikipedia.org/wiki/JPEG#Discrete_cosine_transform"""
result=numpy.zeros(X.shape)
N1,N2=X.shape
def alpha(n):
"""Normalizing function, not sure if necessary"""
if n==0:return 0.353553390593 #sqrt(1/8.)
else:return .5 #sqrt(2/8.)
for (k1,k2),_ in numpy.ndenumerate(X):
sub_result=0.
for n1 in range(N1):
for n2 in range(N2):
sub_result+=X[n1,n2] * cos((pi/N1)*(n1+.5)*k1)*cos((pi/N2)*(n2+.5)*k2)
result[k1,k2]=alpha(k1)*alpha(k2)*sub_result
return result
Looking into something similar myself ATM. Have you looked at SciPy? They have a fft module that I’ve been meaning to look at.
That said you can speend the above up my taking some calculations out of the loop. pi/N1 and pi/N2 should be pre-calulated out side the loop somewhere.
Simon, I don’t think the SciPy fft module includes DCT. Please let me know if you find out otherwise.
Taking some of the calcultions out of the loop is a good idea.
You can precompute most of the two for-loops like this
precomputed = [(n1,n2, pi/N1*(n1+.5), pi/N2*(n2+0.5))
for n1 in range(N1) for n2 in range(N2)]
and then your main loop is more like
for (k1,k2),_ in numpy.ndenumerate(X):
sub_result=0.
for n1, n2, a1, a2 in precomputed:
sub_result+=X[n1,n2] * cos(a1*k1) * cos(a2*k2)
result[k1,k2]=alpha(k1)*alpha(k2)*sub_result
you might even find that
sub_result = sum(X[n1,n2]*cos(a1*k1)*cos(a2*k2) for (n1,n2,a1,a2) in precomputed)
is faster, because the summation is done in C. But since you’re using numpy, you should probably think about using an array for the ‘precomputed’ data, an array for each of the cosine terms, then element-wise product of the three arrays and a summation over all elements.
Good suggestions, Andrew. I think those would help. Maybe I’ll give them a try later.