api_server.py

from flask import Flask, request, jsonify, send_file
from flask_cors import CORS
import os
import sys
import cv2
import numpy as np
import json
import base64
import time
import faiss
import torch
from pathlib import Path
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from datetime import datetime

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from src.backbones import get_model
from src.post_eval import mean_top1p
from src.utils import dists2map

app = Flask(__name__)
CORS(app)

MODEL = None
KNN_INDEX = None
CONFIG = None
MODEL_LOADED = False

UPLOAD_FOLDER = 'api_uploads'
RESULT_FOLDER = 'api_results'
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(RESULT_FOLDER, exist_ok=True)

DEFAULT_RESULTS_DIR = 'results_MVTec/dinov2_vitl14_518/16-shot_preprocess=agnostic_res518'
DEFAULT_OBJECT = 'screw'

def load_model_once():
    global MODEL, KNN_INDEX, CONFIG, MODEL_LOADED

    if MODEL_LOADED:
        print("✓ 使用已加载的模型")
        return MODEL, KNN_INDEX, CONFIG

    print("=" * 60)
    print("正在初始化 AnomalyDINO 模型...")
    print("=" * 60)

    try:
        import yaml

        args_path = os.path.join(DEFAULT_RESULTS_DIR, 'args.yaml')
        if not os.path.exists(args_path):
            raise FileNotFoundError(f"找不到训练配置: {args_path}")

        with open(args_path, 'r') as f:
            CONFIG = yaml.safe_load(f)

        print(f"✓ 加载配置: {args_path}")
        print(f"  - 模型: {CONFIG['model_name']}")
        print(f"  - 分辨率: {CONFIG['resolution']}")
        print(f"  - 数据集: {CONFIG['data_root']}")

        print(f"\n正在加载 DINOv2 模型...")
        MODEL = get_model(
            CONFIG['model_name'],
            'cuda' if torch.cuda.is_available() else 'cpu',
            smaller_edge_size=CONFIG['resolution']
        )
        print("✓ DINOv2 模型加载完成")

        data_root = CONFIG['data_root']
        img_ref_folder = f"{data_root}/{DEFAULT_OBJECT}/train/good/"

        if not os.path.exists(img_ref_folder):
            raise FileNotFoundError(f"找不到参考样本目录: {img_ref_folder}")

        print(f"\n正在加载参考样本: {img_ref_folder}")

        features_ref = []
        img_ref_samples = sorted([f for f in os.listdir(img_ref_folder)
                                 if f.lower().endswith(('.png', '.jpg', '.jpeg'))])

        print(f"找到 {len(img_ref_samples)} 个参考样本")

        with torch.inference_mode():
            for idx, img_ref_n in enumerate(img_ref_samples, 1):
                img_ref_path = os.path.join(img_ref_folder, img_ref_n)
                image_ref = cv2.cvtColor(cv2.imread(img_ref_path, cv2.IMREAD_COLOR),
                                        cv2.COLOR_BGR2RGB)

                image_ref_tensor, grid_size = MODEL.prepare_image(image_ref)
                features_ref_i = MODEL.extract_features(image_ref_tensor)

                mask_ref = MODEL.compute_background_mask(
                    features_ref_i, grid_size,
                    threshold=10, masking_type=True
                )
                features_ref.append(features_ref_i[mask_ref])

                if idx % 5 == 0:
                    print(f"  处理进度: {idx}/{len(img_ref_samples)}")

            features_ref = np.concatenate(features_ref, axis=0).astype('float32')

            print(f"\n正在创建 FAISS 索引...")
            if CONFIG.get('faiss_on_cpu', True):
                KNN_INDEX = faiss.IndexFlatL2(features_ref.shape[1])
            else:
                res = faiss.StandardGpuResources()
                KNN_INDEX = faiss.GpuIndexFlatL2(res, features_ref.shape[1])

            faiss.normalize_L2(features_ref)
            KNN_INDEX.add(features_ref)

            print(f" FAISS 索引创建完成: {features_ref.shape[0]} 个特征向量")

        MODEL_LOADED = True

        print("\n" + "=" * 60)
        print(" 模型初始化完成，准备接收检测请求")
        print("=" * 60 + "\n")

        return MODEL, KNN_INDEX, CONFIG

    except Exception as e:
        print(f" 模型加载失败: {str(e)}")
        import traceback
        traceback.print_exc()
        raise

def detect_image_with_visualization(image_path, output_dir):

    model, knn_index, config = load_model_once()

    start_time = time.time()

    if not os.path.exists(image_path):
        raise FileNotFoundError(f"找不到图片: {image_path}")

    image_test = cv2.cvtColor(cv2.imread(image_path, cv2.IMREAD_COLOR),
                              cv2.COLOR_BGR2RGB)

    with torch.inference_mode():
        image_tensor, grid_size = model.prepare_image(image_test)
        features = model.extract_features(image_tensor)

        mask = model.compute_background_mask(
            features, grid_size,
            threshold=10, masking_type=True
        )
        features_masked = features[mask]

        faiss.normalize_L2(features_masked)

        k_neighbors = config.get('k_neighbors', 1)
        distances, _ = knn_index.search(features_masked, k=k_neighbors)

        if k_neighbors > 1:
            distances = distances.mean(axis=1)

        distances = distances / 2

        anomaly_score = mean_top1p(distances.flatten())

        output_distances = np.zeros_like(mask, dtype=float)
        output_distances[mask] = distances.squeeze()
        d_masked = output_distances.reshape(grid_size)

        anomaly_map = dists2map(d_masked, image_test.shape)

        threshold = 0.38
        is_anomaly = anomaly_score > threshold

        if anomaly_score < 0.15:
            defect_type = "正常"
            confidence = "高"
        elif anomaly_score < threshold:
            defect_type = "正常"
            confidence = "中"
        elif anomaly_score < 0.50:
            defect_type = "轻微异常"
            confidence = "中"
        else:
            defect_type = "严重异常"
            confidence = "高"

        fig, axes = plt.subplots(1, 4, figsize=(20, 5))

        axes[0].imshow(image_test)
        axes[0].set_title('Test', fontsize=12)
        axes[0].axis('off')

        mask_vis = mask.reshape(grid_size[0], grid_size[1]).astype(float)
        mask_vis_resized = cv2.resize(mask_vis, (image_test.shape[1], image_test.shape[0]))

        pca_vis = d_masked.copy()
        pca_vis_resized = cv2.resize(pca_vis, (image_test.shape[1], image_test.shape[0]))

        axes[1].imshow(pca_vis_resized, cmap='rainbow')
        axes[1].imshow(mask_vis_resized, cmap='gray', alpha=0.3)
        axes[1].set_title('Test (PCA + Mask)', fontsize=12)
        axes[1].axis('off')

        im3 = axes[2].imshow(anomaly_map, cmap='jet')
        axes[2].set_title('Patch Distances (1NN)', fontsize=12)
        axes[2].axis('off')

        cbar3 = plt.colorbar(im3, ax=axes[2], fraction=0.046)
        cbar3.ax.tick_params(labelsize=8)

        axes[3].hist(distances.flatten(), bins=50, color='steelblue', edgecolor='black')
        axes[3].axvline(x=threshold, color='red', linestyle='--', linewidth=2,
                       label=f'Threshold: {threshold:.2f}')
        axes[3].set_title('Hist of Distances', fontsize=12)
        axes[3].set_xlabel('Distance', fontsize=10)
        axes[3].set_ylabel('Frequency', fontsize=10)
        axes[3].legend(fontsize=9)
        axes[3].grid(True, alpha=0.3)

        plt.tight_layout()

        os.makedirs(output_dir, exist_ok=True)
        base_name = os.path.splitext(os.path.basename(image_path))[0]
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        result_filename = f'{base_name}_{timestamp}_result.png'
        result_path = os.path.join(output_dir, result_filename)

        plt.savefig(result_path, dpi=150, bbox_inches='tight')
        plt.close()

        fig2, ax = plt.subplots(figsize=(5, 5))
        ax.imshow(image_test)
        ax.axis('off')
        img1_path = os.path.join(output_dir, f'{base_name}_{timestamp}_original.png')
        plt.savefig(img1_path, dpi=150, bbox_inches='tight', pad_inches=0)
        plt.close()

        fig2, ax = plt.subplots(figsize=(5, 5))
        ax.imshow(pca_vis_resized, cmap='rainbow')
        ax.imshow(mask_vis_resized, cmap='gray', alpha=0.3)
        ax.axis('off')
        img2_path = os.path.join(output_dir, f'{base_name}_{timestamp}_pca_mask.png')
        plt.savefig(img2_path, dpi=150, bbox_inches='tight', pad_inches=0)
        plt.close()

        fig2, ax = plt.subplots(figsize=(5, 5))
        ax.imshow(anomaly_map, cmap='jet')
        ax.axis('off')
        img3_path = os.path.join(output_dir, f'{base_name}_{timestamp}_heatmap.png')
        plt.savefig(img3_path, dpi=150, bbox_inches='tight', pad_inches=0)
        plt.close()

        fig2, ax = plt.subplots(figsize=(5, 5))
        ax.hist(distances.flatten(), bins=50, color='steelblue', edgecolor='black')
        ax.axvline(x=threshold, color='red', linestyle='--', linewidth=2)
        ax.grid(True, alpha=0.3)
        img4_path = os.path.join(output_dir, f'{base_name}_{timestamp}_histogram.png')
        plt.savefig(img4_path, dpi=150, bbox_inches='tight')
        plt.close()

        processing_time = time.time() - start_time

        result = {
            'success': True,
            'anomaly_score': float(anomaly_score),
            'is_anomaly': bool(is_anomaly),
            'threshold': float(threshold),
            'defect_type': defect_type,
            'confidence': confidence,
            'processing_time': float(processing_time),
            'num_patches': int(mask.sum()),
            'total_patches': int(mask.size),
            'result_image': result_filename,
            'original_image': f'{base_name}_{timestamp}_original.png',
            'pca_mask_image': f'{base_name}_{timestamp}_pca_mask.png',
            'heatmap_image': f'{base_name}_{timestamp}_heatmap.png',
            'histogram_image': f'{base_name}_{timestamp}_histogram.png',
        }

        return result

@app.route('/health', methods=['GET'])
def health_check():
    """健康检查接口"""
    return jsonify({
        'status': 'ok',
        'model_loaded': MODEL_LOADED,
        'timestamp': datetime.now().isoformat()
    })

@app.route('/detect', methods=['POST'])
def detect():
    try:
        if 'image' not in request.files:
            return jsonify({
                'success': False,
                'error': '没有上传图片'
            }), 400

        file = request.files['image']

        if file.filename == '':
            return jsonify({
                'success': False,
                'error': '文件名为空'
            }), 400

        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"{timestamp}_{file.filename}"
        filepath = os.path.join(UPLOAD_FOLDER, filename)
        file.save(filepath)

        print(f"\n收到检测请求: {filename}")

        result = detect_image_with_visualization(filepath, RESULT_FOLDER)

        print(f"检测完成: 异常得分={result['anomaly_score']:.4f}, "
              f"判定={'异常' if result['is_anomaly'] else '正常'}")

        return jsonify(result)

    except Exception as e:
        print(f" 检测失败: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500

@app.route('/result/<filename>', methods=['GET'])
def get_result_image(filename):
    try:
        file_path = os.path.join(RESULT_FOLDER, filename)
        if os.path.exists(file_path):
            return send_file(file_path, mimetype='image/png')
        else:
            return jsonify({'error': '文件不存在'}), 404
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/model/info', methods=['GET'])
def model_info():
    if MODEL_LOADED:
        return jsonify({
            'loaded': True,
            'config': CONFIG,
            'model_name': CONFIG.get('model_name'),
            'resolution': CONFIG.get('resolution'),
            'object_type': DEFAULT_OBJECT
        })
    else:
        return jsonify({
            'loaded': False,
            'message': '模型未加载'
        })

if __name__ == '__main__':
    print("\n" + "=" * 70)
    print("AnomalyDINO Flask API 服务器")
    print("=" * 70)
    print(f"监听地址: http://localhost:5001")
    print(f"检测对象: {DEFAULT_OBJECT}")
    print(f"模型配置: {DEFAULT_RESULTS_DIR}")
    print("=" * 70 + "\n")

    try:
        load_model_once()
    except Exception as e:
        print(f"️模型预加载失败，将在首次请求时加载")
        print(f"   错误: {str(e)}\n")

    app.run(host='0.0.0.0', port=5001, debug=False, threaded=True)