import numpy as npimport cv2import torchimport torch.nn.functional as Ffrom scipy import ndimagefrom scipy.linalg import sqrtmfrom skimage.metrics import structural_similarityimport warningsimport torchvision.models as modelsimport torchvision.transforms as transformswarnings.filterwarnings('ignore')class ImageMetrics:    def __init__(self, data_range=255):        self.data_range = data_range    def calculate_mse(self, img1, img2):        if img1.shape != img2.shape:            raise ValueError("输入图像形状不一致")        mse = np.mean((img1.astype(np.float64) - img2.astype(np.float64)) ** 2)        return mse    def calculate_psnr(self, img1, img2):        mse = self.calculate_mse(img1, img2)        if mse == 0:            return float('inf')        psnr = 10 * np.log10(self.data_range ** 2 / mse)        return psnr    def calculate_ssim(self, img1, img2, window_size=11, gaussian_weights=True,                       sigma=1.5, use_sample_covariance=True):        # 确保图像是2D或3D        if len(img1.shape) == 2:            img1 = img1[..., np.newaxis]            img2 = img2[..., np.newaxis]        # 计算SSIM        ssim_value, ssim_map = structural_similarity(            img1, img2,            win_size=window_size,            gaussian_weights=gaussian_weights,            sigma=sigma,            use_sample_covariance=use_sample_covariance,            data_range=self.data_range,            channel_axis=2 if img1.shape[2] > 1 else None,            full=True        )        return ssim_value, ssim_map    def calculate_ms_ssim(self, img1, img2, weights=None):        if weights is None:            weights = [0.0448, 0.2856, 0.3001, 0.2363, 0.1333]        levels = len(weights)        mssim = []        mcs = []        # 转换为float        img1 = img1.astype(np.float64)        img2 = img2.astype(np.float64)        for i in range(levels):            if i > 0:                # 下采样                img1 = cv2.pyrDown(img1)                img2 = cv2.pyrDown(img2)            ssim_val, ssim_map = self.calculate_ssim(img1, img2)            if i < levels - 1:                # 计算对比度和结构部分                C1 = (0.01 * self.data_range) ** 2                C2 = (0.03 * self.data_range) ** 2                # 计算局部均值和方差                mu1 = cv2.GaussianBlur(img1, (11, 11), 1.5)                mu2 = cv2.GaussianBlur(img2, (11, 11), 1.5)                sigma1_sq = cv2.GaussianBlur(img1 ** 2, (11, 11), 1.5) - mu1 ** 2                sigma2_sq = cv2.GaussianBlur(img2 ** 2, (11, 11), 1.5) - mu2 ** 2                sigma12 = cv2.GaussianBlur(img1 * img2, (11, 11), 1.5) - mu1 * mu2                # 对比度比较                cs_map = (2 * sigma12 + C2) / (sigma1_sq + sigma2_sq + C2)                mcs.append(np.mean(cs_map))            mssim.append(ssim_val)        # 组合各尺度结果        ms_ssim = np.prod(np.array(mcs) ** np.array(weights[:-1])) * (mssim[-1] ** weights[-1])        return ms_ssim    def calculate_fsim(self, img1, img2):        # 转换为灰度        if len(img1.shape) == 3:            img1_gray = cv2.cvtColor(img1, cv2.COLOR_RGB2GRAY)            img2_gray = cv2.cvtColor(img2, cv2.COLOR_RGB2GRAY)        else:            img1_gray = img1            img2_gray = img2        # 计算相位一致性（简化：使用Sobel梯度）        sobel_x = cv2.Sobel(img1_gray, cv2.CV_64F, 1, 0, ksize=3)        sobel_y = cv2.Sobel(img1_gray, cv2.CV_64F, 0, 1, ksize=3)        pc1 = np.sqrt(sobel_x ** 2 + sobel_y ** 2)        sobel_x = cv2.Sobel(img2_gray, cv2.CV_64F, 1, 0, ksize=3)        sobel_y = cv2.Sobel(img2_gray, cv2.CV_64F, 0, 1, ksize=3)        pc2 = np.sqrt(sobel_x ** 2 + sobel_y ** 2)        # 计算梯度幅度        G1 = pc1.copy()        G2 = pc2.copy()        # 计算相似性        T1 = 0.85        T2 = 160        S_L = (2 * G1 * G2 + T1) / (G1 ** 2 + G2 ** 2 + T1)        S_PC = (2 * pc1 * pc2 + T2) / (pc1 ** 2 + pc2 ** 2 + T2)        PC_m = np.maximum(pc1, pc2)        # 计算FSIM        numerator = np.sum(S_L * S_PC * PC_m)        denominator = np.sum(PC_m)        if denominator == 0:            return 0        fsim = numerator / denominator        return fsimclass NoReferenceMetrics:    def __init__(self):        pass    def calculate_brisque_features(self, img):        if len(img.shape) == 3:            img_gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)        else:            img_gray = img        # 局部归一化        mu = cv2.GaussianBlur(img_gray, (7, 7), 7 / 6)        mu_sq = mu * mu        sigma = np.sqrt(np.abs(cv2.GaussianBlur(img_gray * img_gray, (7, 7), 7 / 6) - mu_sq))        # MSCN系数        mscn = (img_gray - mu) / (sigma + 1)        # 计算统计特征        features = []        # 1. MSCN系数的AGGD参数        from scipy.stats import skew, kurtosis        features.extend([            np.mean(mscn),            np.std(mscn),            skew(mscn.flatten()),            kurtosis(mscn.flatten())        ])        # 2. 水平方向差值        h_diff = mscn[:, :-1] - mscn[:, 1:]        features.extend([            np.mean(h_diff),            np.std(h_diff),            skew(h_diff.flatten()),            kurtosis(h_diff.flatten())        ])        # 3. 垂直方向差值        v_diff = mscn[:-1, :] - mscn[1:, :]        features.extend([            np.mean(v_diff),            np.std(v_diff),            skew(v_diff.flatten()),            kurtosis(v_diff.flatten())        ])        return np.array(features)    def calculate_niqe_features(self, img):        if len(img.shape) == 3:            img_gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)        else:            img_gray = img        # 计算MSCN系数（与BRISQUE相同）        mu = cv2.GaussianBlur(img_gray, (7, 7), 7 / 6)        mu_sq = mu * mu        sigma = np.sqrt(np.abs(cv2.GaussianBlur(img_gray * img_gray, (7, 7), 7 / 6) - mu_sq))        mscn = (img_gray - mu) / (sigma + 1)        # 提取统计特征        from scipy.stats import skew, kurtosis        features = [            np.mean(mscn),            np.var(mscn),            skew(mscn.flatten()),            kurtosis(mscn.flatten())        ]        # 添加其他统计特征        features.append(np.percentile(mscn, 10))        features.append(np.percentile(mscn, 25))        features.append(np.percentile(mscn, 50))        features.append(np.percentile(mscn, 75))        features.append(np.percentile(mscn, 90))        return np.array(features)class DeepLearningMetrics:    """深度学习相关质量指标"""    def __init__(self, device='cuda'if torch.cuda.is_available() else'cpu'):        self.device = torch.device(device)        # 加载预训练模型用于LPIPS        self.vgg = models.vgg16(pretrained=True).features.to(self.device).eval()        for param in self.vgg.parameters():            param.requires_grad = False        # 图像转换        self.transform = transforms.Compose([            transforms.ToTensor(),            transforms.Normalize(mean=[0.485, 0.456, 0.406],                                 std=[0.229, 0.224, 0.225])        ])    def calculate_lpips(self, img1, img2, use_vgg=True):        # 转换图像格式        if len(img1.shape) == 2:            img1 = np.stack([img1, img1, img1], axis=2)            img2 = np.stack([img2, img2, img2], axis=2)        # 转换为tensor        img1_tensor = self.transform(img1).unsqueeze(0).to(self.device)        img2_tensor = self.transform(img2).unsqueeze(0).to(self.device)        if use_vgg:            # 提取VGG特征            features1 = self._extract_vgg_features(img1_tensor)            features2 = self._extract_vgg_features(img2_tensor)            # 计算特征距离（加权）            lpips_value = 0            weights = [1.0, 0.5, 0.25, 0.125, 0.0625]  # 简化权重            for i, (f1, f2) in enumerate(zip(features1, features2)):                if i >= len(weights):                    break                # 归一化特征                f1_norm = F.normalize(f1, p=2, dim=1)                f2_norm = F.normalize(f2, p=2, dim=1)                # 计算L2距离                diff = f1_norm - f2_norm                dist = torch.mean(diff ** 2)                lpips_value += weights[i] * dist.item()        else:            # 使用像素级差异作为简化版本            diff = img1_tensor - img2_tensor            lpips_value = torch.mean(diff ** 2).item()        return lpips_value    def _extract_vgg_features(self, x):        """        提取VGG特征        """        features = []        for i, layer in enumerate(self.vgg):            x = layer(x)            if i in [3, 8, 15, 22, 29]:  # 特定卷积层后                features.append(x)                if len(features) >= 5:  # 只取前5层                    break        return features    def calculate_fid(self, real_images, fake_images, batch_size=64):        # 加载Inception模型        inception = models.inception_v3(pretrained=True, transform_input=False)        inception.fc = torch.nn.Identity()  # 移除分类层        inception = inception.to(self.device).eval()        def get_features(images):            """提取特征"""            features = []            for i in range(0, len(images), batch_size):                batch = images[i:i + batch_size]                batch_tensors = []                for img in batch:                    if len(img.shape) == 2:                        img = np.stack([img, img, img], axis=2)                    img_tensor = self.transform(img).unsqueeze(0)                    batch_tensors.append(img_tensor)                batch_tensor = torch.cat(batch_tensors, 0).to(self.device)                with torch.no_grad():                    feat = inception(batch_tensor)                    features.append(feat.cpu().numpy())            return np.concatenate(features, axis=0)        # 提取特征        real_features = get_features(real_images)        fake_features = get_features(fake_images)        # 计算统计量        mu1, sigma1 = np.mean(real_features, axis=0), np.cov(real_features, rowvar=False)        mu2, sigma2 = np.mean(fake_features, axis=0), np.cov(fake_features, rowvar=False)        # 计算FID        diff = mu1 - mu2        covmean, _ = sqrtm(sigma1.dot(sigma2), disp=False)        if np.iscomplexobj(covmean):            covmean = covmean.real        fid = diff.dot(diff) + np.trace(sigma1 + sigma2 - 2 * covmean)        return fiddef main(img1, img2):    # 确保图像大小相同    if img1.shape != img2.shape:        # 调整大小        img2 = cv2.resize(img2, (img1.shape[1], img1.shape[0]))    # 转换为RGB    img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2RGB)    img2 = cv2.cvtColor(img2, cv2.COLOR_BGR2RGB)    # ===============================    # 1. 基础指标计算    # ===============================    metrics = ImageMetrics(data_range=255)    # MSE    mse = metrics.calculate_mse(img1, img2)    print(f"MSE: {mse:.4f}")    # PSNR    psnr = metrics.calculate_psnr(img1, img2)    print(f"PSNR: {psnr:.2f} dB")    # SSIM    ssim_value, ssim_map = metrics.calculate_ssim(img1, img2)    print(f"SSIM: {ssim_value:.4f}")    # MS-SSIM    ms_ssim = metrics.calculate_ms_ssim(img1, img2)    print(f"MS-SSIM: {ms_ssim:.4f}")    # FSIM    fsim = metrics.calculate_fsim(img1, img2)    print(f"FSIM: {fsim:.4f}")    # ===============================    # 2. 无参考指标计算    # ===============================    no_ref_metrics = NoReferenceMetrics()    niqe = metrics.calculate_niqe(img1, img2)    print(f"NIQE: {niqe:.4f}")    # BRISQUE特征    brisque_features1 = no_ref_metrics.calculate_brisque_features(img1)    brisque_features2 = no_ref_metrics.calculate_brisque_features(img2)    print(f"BRISQUE特征差异: {np.mean(np.abs(brisque_features1 - brisque_features2)):.4f}")    # ===============================    # 3. 深度学习指标    # ===============================    dl_metrics = DeepLearningMetrics(device='cpu')    # LPIPS    lpips_value = dl_metrics.calculate_lpips(img1, img2)    print(f"LPIPS: {lpips_value:.4f}")    return {        'mse': mse,        'psnr': psnr,        'ssim': ssim_value,        'ms_ssim': ms_ssim,        'fsim': fsim,        'niqe': niqe,        'brisque': np.mean(np.abs(brisque_features1 - brisque_features2)),        'lpips': lpips_value,    }if __name__ == "__main__":    # 单对图像评估    # 加载图像    img_path1 = 'reference.jpg'  # 参考图像    img_path2 = 'restored.jpg'  # 复原图像    img1 = cv2.imread(img_path1)    img2 = cv2.imread(img_path2)    results = main(img1, img2)