import cv2
import numpy as np

# 读取视频文件
cap = cv2.VideoCapture('bike.mp4')

# 定义保留振幅的压缩比例列表
keep_list = [1, 0.5, 0.05]

# 循环播放视频文件，同时显示灰度视频、原振幅谱和平移后的振幅谱
while cap.isOpened():
    # 逐帧读取视频，ret 为布尔值，表示是否成功读取帧，frame 为当前帧的图像数据
    ret, frame = cap.read()
    if not ret:
        print("没有内容，退出啦 :) ")
        break
    # 使用 cv2.cvtColor() 将当前帧的彩色图像转换为灰度图
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

    # 对灰度图进行二维快速傅里叶变换，计算振幅谱并放大
    fft = np.fft.fft2(gray)
    mag_spec = 20 * np.log(np.abs(fft))

    mag_spec_sorted = np.sort(mag_spec.ravel())  # 对振幅谱进行排序

    # 循环处理不同压缩比例
    for i, keep in enumerate(keep_list):
        # 计算保留振幅的阈值
        thresh = mag_spec_sorted[int((1 - keep) * gray.size)] # 计算位置的振幅阈值

        # 创建掩码，保留振幅大于阈值的频率分量
        mask = mag_spec > thresh

        # 使用掩码压缩傅里叶系数（模拟通过网络传输压缩后的系数）
        fCoef_compressed = fft * mask

        # 使用压缩后的傅里叶系数进行逆傅里叶变换，得到压缩后的影响
        img_cp = np.fft.ifft2(fCoef_compressed).real

        # 将灰度图像转换为 8 位无符号整型（0-255）
        img_cp = np.uint8(img_cp)
        # 在相应窗口中显示压缩后的视频
        if i == 0:  # 当i==0时，在名为 "original" 的窗口中显示原始影像
            cv2.imshow('original', gray)
        elif i == 1:    # i==1时，在名为 "0.5" 的窗口中显示压缩比为0.5影像
            cv2.imshow('0.5', img_cp)
        elif i == 2:   # i==2时，在名为 "0.05" 的窗口中显示压缩比为0.05影像
            cv2.imshow('0.05', img_cp)

    # 每隔 56ms 检查一次用户输入，如果按下 'q' 键，退出循环
    if cv2.waitKey(56) & 0xFF == ord('q'):
        break

# 释放当前视频捕捉资源
cap.release()

# 关闭所有的 GUI 窗口
cv2.destroyAllWindows()