import cv2
from skimage.metrics import structural_similarity as ssim
import numpy as np

class VideoSSIM:
    def __init__(self, video_path):
        # 初始化视频路径和打开视频
        self.video_path = video_path
        self.cap = cv2.VideoCapture(video_path)
        if not self.cap.isOpened():
            raise ValueError(f"Error opening video file: {video_path}")
        
        # 获取视频帧率（fps）
        self.fps = self.cap.get(cv2.CAP_PROP_FPS)
        
        # 获取视频的尺寸
        self.frame_width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        self.frame_height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    def calculate_ssim(self):
        # 读取视频的第一帧
        ret, prev_frame = self.cap.read()
        if not ret:
            raise ValueError("Error reading the first frame")

        prev_frame_gray = cv2.cvtColor(prev_frame, cv2.COLOR_BGR2GRAY)  # 转换为灰度图
        
        # 存储相似性数值
        ssim_values = []
        
        # 从第二帧开始逐帧计算与前一帧的结构相似性
        frame_index = 2  # 从第二帧开始
        while True:
            ret, curr_frame = self.cap.read()
            if not ret:
                break
            
            # 转换当前帧为灰度图
            curr_frame_gray = cv2.cvtColor(curr_frame, cv2.COLOR_BGR2GRAY)
            
            # 计算SSIM
            ssim_value, _ = ssim(prev_frame_gray, curr_frame_gray, full=True)
            ssim_values.append((frame_index, ssim_value))  # 记录帧索引和SSIM值
            
            # 更新前一帧
            prev_frame_gray = curr_frame_gray
            frame_index += 1
        
        # 释放视频捕捉对象
        self.cap.release()
        
        return ssim_values

    def save_ssim_to_file(self, output_file="ssim_values.txt"):
        ssim_values = self.calculate_ssim()
        
        # 保存SSIM值到文件
        with open(output_file, "w") as f:
            for frame_idx, ssim_value in ssim_values:
                f.write(f"Frame {frame_idx}: SSIM = {ssim_value}\n")
        
        print(f"SSIM values saved to {output_file}")

# 使用示例
video_ssim = VideoSSIM("demo.mp4")  # 替换成你的视频文件路径
video_ssim.save_ssim_to_file("demo_ssim_output.txt")