2D DCT

2d_dct.py

from math import sqrt, cos, pi
import numpy as np 
import matplotlib.pyplot as plt

M=N=8

def alpha_p(p: int) -> float:
    if p == 0:
        return 1/sqrt(M)
    else:
        return sqrt(2/M)

def alpha_q(q: int) -> float:
    if q == 0:
        return 1/sqrt(N)
    else:
        return sqrt(2/N)

def main():
    W = []
    for p in range(M):
        W_q = []
        for q in range(N):
            img = []
            for m in range(M):
                row = []
                for n in range(N):
                    row.append(alpha_p(p) * alpha_q(q) * cos(pi*(2*m+1)*p/(2*M)) * cos(pi*(2*n+1)*q/(2*N)))
                img.append(row)
            W_q.append(img)
        W.append(W_q)

    fig, ax = plt.subplots(M, N)
    for m in range(M):
        for n in range(N):
            ax[m][n].imshow(W[m][n], cmap="gray")
            ax[m][n].axis("off")
    plt.show()

if __name__ == "__main__":
    main()