计算BRISQUE分数

使用BRISQUE分数来评判超分的效果

import pyiqa
import cv2
import numpy as np
import torch 
import sys

# 创建 BRISQUE 模型
iqa_metric = pyiqa.create_metric('brisque', device='cpu')

def calculate_video_brisque(video_path, sample_interval=30):
    scores = []
    cap = cv2.VideoCapture(video_path)
    frame_count = 0
    
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        # 每隔 sample_interval 帧采样一次，平衡速度与准确性
        if frame_count % sample_interval == 0:
            # pyiqa 要求输入是 RGB 格式且范围为 [0, 1]
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frame_tensor = torch.from_numpy(frame_rgb).float() / 255.0
            # 调整维度顺序为 (C, H, W)
            frame_tensor = frame_tensor.permute(2, 0, 1).unsqueeze(0)
            
            score = iqa_metric(frame_tensor)
            scores.append(score.item())
        frame_count += 1
        
    cap.release()
    return np.mean(scores) if scores else None

def process_video2(filename):
    score = calculate_video_brisque(filename)
    
    if score is None:
        print(f"错误：无法计算 {filename} 的分数")
        sys.exit(1)   # 程序终止，返回状态码 1
    
    print(f"{filename} 的平均 BRISQUE 分数: {score:.4f}")
    return score


process_video2("e:/1.ts")
sys.exit()