import cv2  # 加载OPENCV模块
import numpy as np  # 加载numpy模块


def grey_scale(image):
    img_gray = image
    rows, cols = img_gray.shape
    flat_gray = img_gray.reshape((cols * rows,)).tolist()
    A = min(flat_gray)
    B = max(flat_gray)
    print('A = %d,B = %d' % (A, B))
    output = np.uint8(255 / (B - A) * (img_gray - A) + 0.5)
    return output


def DCT(img):  # DCT函数
    Gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)  # 以灰度模式读入图像
    row, height = Gray_img.shape  # 读取图像的大小
    F = cv2.dct(np.float32(Gray_img))  # 将图像转为float32类型，并进行DCT变换
    Fimg = np.uint8(np.abs(F))  # 取绝对值得到幅度谱，将转换为uint8类型方便显示
    cf = 60  # 截止频率
    mask = np.zeros((row, height), np.uint8)
    # 设计二值掩膜将DCT域的高频部分消除
    mask[0:cf, 0:cf] = 1
    dct = F * mask
    # 压缩后的DCT域图像
    dct_img = cv2.idct(dct)
    F1 = cv2.dct(np.float32(dct_img))  # 将图像转为float32类型，并进行DCT变换
    F1img = np.uint8(np.abs(F1))  # 取绝对值得到幅度谱，将转换为uint8类型方便显示
    cv2.imshow('img', img)  # 显示原图像
    cv2.imshow('Fimg', grey_scale(Fimg))  # 显示DCT域图像
    cv2.imshow('dct_img', grey_scale(dct_img))  # 显示压缩后的DCT域图像
    cv2.imshow('F1img', grey_scale(F1img))  # 显示重建后的图像
    cv2.waitKey(0)  # 等待时间
    cv2.destroyAllWindows()  # 关闭所有创建的窗口


if __name__ == '__main__':  # 主函数
    img = cv2.imread('desert.jpg')  # 读取图像
    DCT(img)  # 调用DCT函数