#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
使用支持向量机(SVM)进行工业缺陷检测
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.svm import SVC, LinearSVC
from sklearn.model_selection import train_test_split, GridSearchCV, cross_val_score
from sklearn.metrics import (accuracy_score, precision_score, recall_score, f1_score,
                            classification_report, confusion_matrix, roc_curve,
                            roc_auc_score, precision_recall_curve)
from sklearn.preprocessing import StandardScaler, MinMaxScaler
from sklearn.decomposition import PCA
from sklearn.pipeline import Pipeline
import os
import warnings
import joblib
from sklearn.feature_extraction.image import extract_patches_2d
from skimage import io, color, feature, filters
import cv2

# 忽略警告
warnings.filterwarnings('ignore')

class IndustrialDefectDetector:
    """
    使用SVM进行工业缺陷检测的类
    """
    def __init__(self, kernel='rbf', C=1.0, gamma='scale', random_state=42):
        """
        初始化工业缺陷检测器
        
        参数:
        kernel: SVM核函数类型，可选'linear', 'poly', 'rbf', 'sigmoid'
        C: 正则化参数
        gamma: 核系数
        random_state: 随机种子
        """
        self.kernel = kernel
        self.C = C
        self.gamma = gamma
        self.random_state = random_state
        self.model = SVC(
            kernel=kernel,
            C=C,
            gamma=gamma,
            probability=True,
            random_state=random_state
        )
        self.scaler = StandardScaler()
        self.feature_names = []
        self.class_names = ['正常', '缺陷']  # 0表示正常，1表示缺陷
        self.is_trained = False
    
    def load_data(self, data_dir=None, image_size=(64, 64), file_path=None):
        """
        加载工业缺陷检测数据
        
        参数:
        data_dir: 数据目录，包含'normal'和'defect'子目录
        image_size: 调整图像大小
        file_path: 特征数据文件路径，如果提供则从文件加载
        
        返回:
        特征X和标签y
        """
        if file_path and os.path.exists(file_path):
            # 从文件加载数据
            try:
                data = pd.read_csv(file_path)
                print(f"从文件加载数据成功: {file_path}")
                X = data.drop('label', axis=1).values
                y = data['label'].values
                self.feature_names = data.columns.tolist()[:-1]
                return X, y
            except Exception as e:
                print(f"从文件加载数据失败: {e}")
                print("尝试从图像数据生成特征...")
        
        if not data_dir or not os.path.exists(data_dir):
            print("未提供有效的数据目录，将生成模拟数据...")
            return self._create_mock_data()
        
        print(f"从{data_dir}加载图像数据...")
        
        # 定义特征和标签列表
        X = []
        y = []
        
        # 加载正常样本
        normal_dir = os.path.join(data_dir, 'normal')
        if os.path.exists(normal_dir):
            for img_file in os.listdir(normal_dir):
                if img_file.endswith(('.jpg', '.jpeg', '.png', '.bmp')):
                    img_path = os.path.join(normal_dir, img_file)
                    features = self._extract_features_from_image(img_path, image_size)
                    X.append(features)
                    y.append(0)  # 正常类别
        
        # 加载缺陷样本
        defect_dir = os.path.join(data_dir, 'defect')
        if os.path.exists(defect_dir):
            for img_file in os.listdir(defect_dir):
                if img_file.endswith(('.jpg', '.jpeg', '.png', '.bmp')):
                    img_path = os.path.join(defect_dir, img_file)
                    features = self._extract_features_from_image(img_path, image_size)
                    X.append(features)
                    y.append(1)  # 缺陷类别
        
        if len(X) == 0:
            print("未找到图像数据，将生成模拟数据...")
            return self._create_mock_data()
        
        # 转换为numpy数组
        X = np.array(X)
        y = np.array(y)
        
        print(f"已加载数据，形状: X={X.shape}, y={y.shape}")
        
        # 生成特征名称
        self.feature_names = [f'feature_{i+1}' for i in range(X.shape[1])]
        
        # 保存到文件
        if file_path:
            try:
                # 创建DataFrame
                data = pd.DataFrame(X, columns=self.feature_names)
                data['label'] = y
                
                # 保存到CSV
                data.to_csv(file_path, index=False)
                print(f"数据已保存到文件: {file_path}")
            except Exception as e:
                print(f"保存数据到文件失败: {e}")
        
        return X, y
    
    def _extract_features_from_image(self, image_path, image_size=(64, 64)):
        """
        从图像中提取特征
        
        参数:
        image_path: 图像文件路径
        image_size: 调整图像大小
        
        返回:
        特征向量
        """
        # 读取图像
        try:
            img = cv2.imread(image_path)
            if img is None:
                raise ValueError(f"无法读取图像: {image_path}")
            
            # 调整图像大小
            img = cv2.resize(img, image_size)
            
            # 转换为灰度图
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            
            # 提取特征
            features = []
            
            # 1. 提取HOG特征
            hog_features = feature.hog(
               # 计算图像的HOG（方向梯度直方图）特征
                # 该函数用于从图像中提取特征，通过计算图像的梯度来了解图像中对象的形状
                gray,
                # orientations: 指定每个细胞单元中计算的梯度方向的数量
                # 一个方向表示一个梯度方向，用于捕捉图像中的纹理信息
                orientations=8,
                # pixels_per_cell: 指定每个细胞单元的大小
                # 细胞单元是图像的小区域，此参数决定了这些区域的大小
                # 这对于确定特征向量的维度非常重要
                pixels_per_cell=(8, 8),
                # cells_per_block: 指定每个块包含多少个细胞单元
                # 块是细胞单元的重叠区域，此参数决定了这些块的大小
                # 它影响到特征向量的维度和性能
                cells_per_block=(2, 2),
                # block_norm: 指定使用的块归一化方法
                # 'L2-Hys' 是一种L2归一化方法，它有助于减少光照变化的影响
                block_norm='L2-Hys',
                # feature_vector: 指定是否将HOG特征返回为特征向量
                # 如果为True，特征将被展平为一个向量，适用于机器学习算法的输入
                feature_vector=True

            )
            # 将HOG特征添加到特征集合中
            features.extend(hog_features)

            
            # 2. 添加一些基本统计特征
            features.append(np.mean(gray))  # 平均灰度
            features.append(np.std(gray))   # 标准差
            features.append(np.max(gray) - np.min(gray))  # 范围
            
            # 3. 提取边缘特征
            edges = filters.sobel(gray)
            features.append(np.mean(edges))
            features.append(np.std(edges))
            
            return np.array(features)
            
        except Exception as e:
            print(f"提取图像特征时出错: {e}")
            # 返回全零特征向量
            return np.zeros(100)
    
    def _create_mock_data(self, n_samples=200, n_features=100):
        """
        创建模拟工业缺陷检测数据
        
        参数:
        n_samples: 样本数量
        n_features: 特征数量
        
        返回:
        特征X和标签y
        """
        print("创建模拟工业缺陷检测数据...")
        
        # 设置随机种子
        np.random.seed(self.random_state)
        
        # 生成正常样本特征
        # 计算正常样本的数量，取总样本数的一半
        n_normal = n_samples // 2

        # 初始化正常样本的均值向量，所有特征的均值均为0
        normal_mean = np.zeros(n_features)

        # 初始化正常样本的协方差矩阵，对角线元素为0.5，表示特征间的协方差为0，即特征相互独立
        normal_cov = np.eye(n_features) * 0.5

        # 生成正常样本数据集，使用多元正态分布
        X_normal = np.random.multivariate_normal(normal_mean, normal_cov, n_normal)

        
        # 生成缺陷样本特征，与正常样本有一些差异
        # 计算缺陷样本的数量
        n_defect = n_samples - n_normal

        # 初始化缺陷样本的特征均值向量，所有特征的均值默认为0
        defect_mean = np.zeros(n_features)

        # 设置前20个特征的均值为2.0，表示这些特征在缺陷样本中有明显不同的表现
        defect_mean[0:20] = 2.0  # 前20个特征有明显不同

        # 初始化缺陷样本的特征协方差矩阵，表示各特征之间相互独立，且每个特征的方差为0.7
        defect_cov = np.eye(n_features) * 0.7

        # 生成缺陷样本的数据集，使用多变量正态分布，均值和协方差矩阵已定义
        X_defect = np.random.multivariate_normal(defect_mean, defect_cov, n_defect)

        
        # 合并样本
       # 将正常样本和缺陷样本的特征矩阵垂直堆叠在一起
        X = np.vstack((X_normal, X_defect))
        # 创建目标变量数组，正常样本标记为0，缺陷样本标记为1
        y = np.hstack((np.zeros(n_normal), np.ones(n_defect)))

        
        # 打乱样本顺序
        indices = np.arange(n_samples)
        np.random.shuffle(indices)
        X = X[indices]
        y = y[indices]
        
        # 生成特征名称
        self.feature_names = [f'feature_{i+1}' for i in range(n_features)]
        
        print(f"已创建模拟数据，形状: X={X.shape}, y={y.shape}")
        return X, y
    
    def prepare_data(self, X, y, test_size=0.2, apply_pca=False, n_components=0.95):
        """
        准备数据，包括缩放和PCA降维（可选）
        
        参数:
        X: 特征数据
        y: 标签
        test_size: 测试集比例
        apply_pca: 是否应用PCA降维
        n_components: PCA保留的方差比例或组件数
        
        返回:
        X_train, X_test, y_train, y_test
        """
        # 数据缩放
        X_scaled = self.scaler.fit_transform(X)
        
        # 可选的PCA降维
        if apply_pca:
            pca = PCA(n_components=n_components)
            X_scaled = pca.fit_transform(X_scaled)
            explained_var = np.sum(pca.explained_variance_ratio_)
            n_components_actual = pca.n_components_
            print(f"PCA降维后的特征数量: {n_components_actual}")
            print(f"保留的方差比例: {explained_var:.4f}")
            
            # 更新特征名称
            self.feature_names = [f'PC_{i+1}' for i in range(n_components_actual)]
        
        # 分割数据集
        X_train, X_test, y_train, y_test = train_test_split(
            X_scaled, y, test_size=test_size, random_state=self.random_state, stratify=y
        )
        
        print(f"训练集大小: {X_train.shape[0]}, 测试集大小: {X_test.shape[0]}")
        unique, counts = np.unique(y, return_counts=True)
        class_dist = dict(zip(unique, counts))
        print(f"类别分布: {class_dist}")
        
        return X_train, X_test, y_train, y_test
    
    def train(self, X_train, y_train):
        """
        训练SVM模型
        
        参数:
        X_train: 训练特征
        y_train: 训练标签
        
        返回:
        训练好的模型
        """
        print("\n开始训练SVM模型...")
        self.model.fit(X_train, y_train)
        self.is_trained = True
        print("SVM模型训练完成")
        
        # 打印模型参数
        print("\nSVM模型参数:")
        for param, value in self.model.get_params().items():
            print(f"{param}: {value}")
        
        # 计算训练集上的性能
        y_train_pred = self.model.predict(X_train)
        train_accuracy = accuracy_score(y_train, y_train_pred)
        print(f"\n训练集准确率: {train_accuracy:.4f}")
        
        # 计算支持向量数量
        if hasattr(self.model, 'support_vectors_'):
            n_support = self.model.n_support_
            print(f"支持向量数量: {sum(n_support)}")
            for i, count in enumerate(n_support):
                print(f"类别 {i} 的支持向量数量: {count}")
        
        return self.model
    
    def evaluate(self, X_test, y_test):
        """
        评估模型性能
        
        参数:
        X_test: 测试特征
        y_test: 测试标签
        
        返回:
        评估指标字典
        """
        if not self.is_trained:
            print("模型尚未训练，请先训练模型")
            return None
        
        # 预测
        y_pred = self.model.predict(X_test)
        y_prob = None
        if hasattr(self.model, "predict_proba"):
            y_prob = self.model.predict_proba(X_test)[:, 1]
        
        # 计算评估指标
        accuracy = accuracy_score(y_test, y_pred)
        precision = precision_score(y_test, y_pred, zero_division=0)
        recall = recall_score(y_test, y_pred, zero_division=0)
        f1 = f1_score(y_test, y_pred, zero_division=0)
        
        # 混淆矩阵
        cm = confusion_matrix(y_test, y_pred)
        
        # ROC曲线和AUC
        roc_auc = None
        if y_prob is not None:
            roc_auc = roc_auc_score(y_test, y_prob)
        
        # 打印评估结果
        print("\n模型评估结果:")
        print(f"准确率: {accuracy:.4f}")
        print(f"精确率: {precision:.4f}")
        print(f"召回率: {recall:.4f}")
        print(f"F1得分: {f1:.4f}")
        if roc_auc is not None:
            print(f"ROC AUC: {roc_auc:.4f}")
        
        print("\n分类报告:")
        print(classification_report(y_test, y_pred, target_names=self.class_names))
        
        print("\n混淆矩阵:")
        print(cm)
        
        # 返回评估指标
        metrics = {
            'accuracy': accuracy,  #准确率
            'precision': precision,  #精确率
            'recall': recall,        #recall
            'f1': f1,               #F1得分
            'roc_auc': roc_auc,
            'confusion_matrix': cm,
            'classification_report': classification_report(y_test, y_pred, target_names=self.class_names, output_dict=True)
        }
        
        return metrics
    
    def plot_confusion_matrix(self, cm, class_names=None, save_path="confusion_matrix.png"):
        """
        绘制混淆矩阵
        
        参数:
        cm: 混淆矩阵
        class_names: 类别名称列表
        save_path: 保存路径
        
        返回:
        无
        """
        if class_names is None:
            class_names = self.class_names
        
        plt.figure(figsize=(8, 6))
        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=class_names, yticklabels=class_names)
        plt.xlabel('预测类别')
        plt.ylabel('真实类别')
        plt.title('混淆矩阵')
        plt.tight_layout()
        plt.savefig(save_path)
        plt.close()
        
        print(f"混淆矩阵图已保存到: {save_path}")
    
    def plot_roc_curve(self, y_test, y_prob, save_path="roc_curve.png"):
        """
        绘制ROC曲线
        
        参数:
        y_test: 真实标签
        y_prob: 预测概率
        save_path: 保存路径
        
        返回:
        无
        """
        fpr, tpr, _ = roc_curve(y_test, y_prob)
        roc_auc = roc_auc_score(y_test, y_prob)
        
        plt.figure(figsize=(8, 6))
        plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC曲线 (AUC = {roc_auc:.3f})')
        plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
        plt.xlim([0.0, 1.0])
        plt.ylim([0.0, 1.05])
        plt.xlabel('假阳性率')
        plt.ylabel('真阳性率')
        plt.title('接收者操作特征曲线')
        plt.legend(loc='lower right')
        plt.tight_layout()
        plt.savefig(save_path)
        plt.close()
        
        print(f"ROC曲线图已保存到: {save_path}")
    
    def optimize_parameters(self, X_train, y_train, param_grid=None, cv=5):
        """
        优化SVM参数
        
        参数:
        X_train: 训练特征
        y_train: 训练标签
        param_grid: 参数网格，如果为None则使用默认网格
        cv: 交叉验证折数
        
        返回:
        最佳参数
        """
        if param_grid is None:
            param_grid = {
                'C': [0.1, 1, 10, 100],
                'gamma': ['scale', 'auto', 0.01, 0.1, 1],
                'kernel': ['linear', 'rbf', 'poly', 'sigmoid']
            }
        
        print("\n开始参数优化...")
        
        # 创建网格搜索
        grid_search = GridSearchCV(
            SVC(probability=True, random_state=self.random_state),
            param_grid,
            cv=cv,
            scoring='f1',
            n_jobs=-1,
            verbose=1
        )
        
        # 执行网格搜索
        grid_search.fit(X_train, y_train)
        
        # 输出最佳参数
        print("\n最佳参数:")
        for param, value in grid_search.best_params_.items():
            print(f"{param}: {value}")
        
        print(f"最佳F1得分: {grid_search.best_score_:.4f}")
        
        # 更新模型参数
        self.model = SVC(
            **grid_search.best_params_,
            probability=True,
            random_state=self.random_state
        )
        
        # 使用最佳参数训练模型
        self.train(X_train, y_train)
        
        return grid_search.best_params_
    
    def save_model(self, model_path="svm_defect_detector.pkl"):
        """
        保存模型
        
        参数:
        model_path: 模型保存路径
        
        返回:
        无
        """
        if not self.is_trained:
            print("模型尚未训练，请先训练模型")
            return
        
        try:
            # 创建包含模型和缩放器的字典
            model_dict = {
                'model': self.model,
                'scaler': self.scaler,
                'feature_names': self.feature_names,
                'class_names': self.class_names
            }
            
            # 保存模型
            joblib.dump(model_dict, model_path)
            print(f"模型已保存到: {model_path}")
        except Exception as e:
            print(f"保存模型失败: {e}")
    
    def load_model(self, model_path="svm_defect_detector.pkl"):
        """
        加载模型
        
        参数:
        model_path: 模型路径
        
        返回:
        加载是否成功
        """
        try:
            # 加载模型
            model_dict = joblib.load(model_path)
            
            # 恢复模型和缩放器
            self.model = model_dict['model']
            self.scaler = model_dict['scaler']
            self.feature_names = model_dict['feature_names']
            self.class_names = model_dict['class_names']
            self.is_trained = True
            
            print(f"模型已从{model_path}加载")
            return True
        except Exception as e:
            print(f"加载模型失败: {e}")
            return False
    
    def predict_defect(self, X, threshold=0.5):
        """
        预测样本是否存在缺陷
        
        参数:
        X: 输入特征
        threshold: 概率阈值，高于该值则判定为缺陷
        
        返回:
        预测结果和概率
        """
        if not self.is_trained:
            print("模型尚未训练，请先训练模型")
            return None, None
        
        # 确保X是2D数组
        if X.ndim == 1:
            X = X.reshape(1, -1)
        
        # 应用特征缩放
        X_scaled = self.scaler.transform(X)
        
        # 预测类别
        y_pred = self.model.predict(X_scaled)
        
        # 预测概率
        y_prob = None
        if hasattr(self.model, "predict_proba"):
            y_prob = self.model.predict_proba(X_scaled)[:, 1]
            # 使用阈值进行决策
            y_pred = (y_prob >= threshold).astype(int)
        
        result_text = []
        for i, (pred, prob) in enumerate(zip(y_pred, y_prob if y_prob is not None else [None] * len(y_pred))):
            status = "缺陷" if pred == 1 else "正常"
            confidence = f", 置信度: {prob:.4f}" if prob is not None else ""
            result_text.append(f"样本 {i+1}: {status}{confidence}")
        
        return y_pred, y_prob, result_text
    
    def detect_defect_in_image(self, image_path, image_size=(64, 64), threshold=0.5):
        """
        检测图像中的缺陷
        
        参数:
        image_path: 图像文件路径
        image_size: 调整图像大小
        threshold: 概率阈值，高于该值则判定为缺陷
        
        返回:
        是否有缺陷，缺陷概率
        """
        if not self.is_trained:
            print("模型尚未训练，请先训练模型")
            return None, None
        
        # 提取图像特征
        features = self._extract_features_from_image(image_path, image_size)
        
        # 预测
        y_pred, y_prob, result_text = self.predict_defect(features, threshold)
        
        return y_pred[0], y_prob[0] if y_prob is not None else None, result_text[0]


def example_defect_detection():
    """工业缺陷检测示例"""
    print("\n=== 工业缺陷检测示例 ===\n")
    
    # 初始化缺陷检测器
    detector = IndustrialDefectDetector(kernel='rbf', C=10.0, gamma='scale')
    
    # 加载数据
    X, y = detector.load_data()
    
    # 准备数据
    X_train, X_test, y_train, y_test = detector.prepare_data(X, y, apply_pca=True, n_components=0.95)
    
    # 训练模型
    detector.train(X_train, y_train)
    
    # 评估模型
    metrics = detector.evaluate(X_test, y_test)
    
    # 绘制混淆矩阵
    detector.plot_confusion_matrix(metrics['confusion_matrix'])
    
    # 如果有概率输出，绘制ROC曲线
    if metrics['roc_auc'] is not None:
        y_prob = detector.model.predict_proba(X_test)[:, 1]
        detector.plot_roc_curve(y_test, y_prob)
    
    # 优化参数（可选）
    # best_params = detector.optimize_parameters(X_train, y_train)
    
    # 保存模型
    detector.save_model()
    
    # 测试预测
    test_idx = np.random.randint(0, len(X_test))
    test_X = X_test[test_idx:test_idx+1]
    test_y = y_test[test_idx]
    
    y_pred, y_prob, result_text = detector.predict_defect(test_X)
    
    print("\n预测示例:")
    print(f"真实标签: {'缺陷' if test_y == 1 else '正常'}")
    print(result_text[0])
    
    return detector


if __name__ == "__main__":
    # 运行工业缺陷检测示例
    detector = example_defect_detection()

-------------------------------------------------下面是调用流程图--------------------------------------------------------+--------------------------------------------------+
| 工业缺陷检测流程开始                             |
+--------------------------------------------------+
            |
            v
+--------------------------------------------------+
| 初始化IndustrialDefectDetector                   |
| (设置SVM参数：核函数、正则化等)                  |
+--------------------------------------------------+
            |
            v
+--------------------------------------------------+
| 加载数据                                         |
| - 从CSV文件加载（如果提供路径）                  |
| - 否则从图像目录提取特征                         |
| - 如果没有有效数据，生成模拟数据                 |
+--------------------------------------------------+
            |
            v
+--------------------------------------------------+
| 特征提取（仅图像数据）                           |
| - HOG特征                                        |
| - 灰度统计特征（均值、标准差、范围）             |
| - 边缘特征                                       |
+--------------------------------------------------+
            |
            v
+--------------------------------------------------+
| 数据预处理                                       |
| - 标准化（StandardScaler）                       |
| - PCA降维（可选）                                |
| - 分割训练集/测试集                              |
+--------------------------------------------------+
            |
            v
+--------------------------------------------------+
| 训练SVM模型                                      |
| - 使用训练集训练                                 |
| - 输出训练时间、准确率、支持向量数量             |
| - 输出模型参数                                   |
+--------------------------------------------------+
            |
            v
+--------------------------------------------------+
| 模型评估                                         |
| - 在测试集上评估                                 |
| - 输出准确率、精确率、召回率、F1分数              |
| - 输出分类报告和混淆矩阵                         |
| - 可选绘制ROC曲线和AUC值                         |
+--------------------------------------------------+
            |
            v
+--------------------------------------------------+
| 参数优化（可选）                                 |
| - 使用GridSearchCV进行网格搜索                   |
| - 选择最佳参数（C、gamma、kernel）               |
| - 重新训练模型                                   |
+--------------------------------------------------+
            |
            v
+--------------------------------------------------+
| 保存/加载模型（持久化）                          |
| - 将训练好的模型保存到磁盘                      |
| - 从磁盘加载已训练好的模型                     |
+--------------------------------------------------+
            |
            v
+--------------------------------------------------+
| 运行预测                                         |
| - 对新样本进行缺陷预测                          |
| - 可设定概率阈值判断是否为缺陷                  |
+--------------------------------------------------+
            |
            v
+--------------------------------------------------+
| 图像缺陷检测（可选）                             |
| - 对单张图像提取特征                            |
| - 预测是否存在缺陷                              |
+--------------------------------------------------+
            |
            v
+--------------------------------------------------+
| 结果可视化                                       |
| - 绘制混淆矩阵                                  |
| - 绘制ROC曲线                                   |
+--------------------------------------------------+
            |
            v
+--------------------------------------------------+
| 工业缺陷检测流程结束                             |
+--------------------------------------------------+