Ten-category_eval.py

import json
import os
from sklearn.metrics import (
    accuracy_score,
    precision_recall_fscore_support,
    confusion_matrix
)
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

# Define classification labels order (must match training)
LABELS = ["docs", "perf", "style", "refactor", "feat",
          "fix", "test", "ci", "build", "chore"]


def load_data(input_file):
    """Load test result data and perform integrity check"""
    with open(input_file, 'r', encoding='utf-8') as f:
        data = json.load(f)

    # Data validation
    required_fields = ["message_type", "Ten-category_result"]
    for idx, item in enumerate(data):
        for field in required_fields:
            if field not in item:
                raise ValueError(f"Record {idx} missing required field: {field}")

    return data


def calculate_metrics(data):
    """Calculate classification metrics"""
    true_labels = []
    pred_labels = []

    for item in data:
        true_labels.append(item["message_type"])
        pred_labels.append(item["Ten-category_result"])

    # Filter invalid predictions
    valid_indices = [i for i, p in enumerate(pred_labels) if p in LABELS]
    filtered_true = [true_labels[i] for i in valid_indices]
    filtered_pred = [pred_labels[i] for i in valid_indices]

    # Calculate various metrics
    accuracy = accuracy_score(filtered_true, filtered_pred)
    precision, recall, f1, support = precision_recall_fscore_support(
        filtered_true, filtered_pred, labels=LABELS, average=None
    )
    macro_p, macro_r, macro_f1, _ = precision_recall_fscore_support(
        filtered_true, filtered_pred, average='macro'
    )

    return {
        "true_labels": filtered_true,
        "pred_labels": filtered_pred,
        "accuracy": accuracy,
        "precision": precision,
        "recall": recall,
        "f1": f1,
        "support": support,
        "macro_p": macro_p,
        "macro_r": macro_r,
        "macro_f1": macro_f1
    }


def generate_confusion_matrix(true_labels, pred_labels, output_path):
    """Generate confusion matrix visualization"""
    plt.figure(figsize=(12, 10))
    cm = confusion_matrix(true_labels, pred_labels, labels=LABELS)

    # Use heatmap style
    sns.heatmap(cm, annot=True, fmt='d', cmap='Reds',
                xticklabels=LABELS, yticklabels=LABELS,
                cbar=False, linewidths=0.5)

    # plt.title('Confusion Matrix', fontsize=14)
    plt.xlabel('Predicted Label', fontsize=20)
    plt.ylabel('True Label', fontsize=20)
    plt.xticks(rotation=45, ha='right')
    plt.yticks(rotation=0)

    plt.tick_params(axis='x', labelsize=20, rotation=45)
    plt.tick_params(axis='y', labelsize=20, rotation=0)

    plt.tight_layout()
    plt.savefig(output_path, dpi=300, bbox_inches='tight')
    plt.close()


def save_metrics_report(metrics, output_path):
    """Save metrics report"""
    # Category metrics
    df = pd.DataFrame({
        "Category": LABELS,
        "Precision": metrics["precision"],
        "Recall": metrics["recall"],
        "F1-Score": metrics["f1"],
        "Support": metrics["support"]
    })

    # Add macro average
    macro_row = pd.DataFrame({
        "Category": ["Macro Avg"],
        "Precision": [metrics["macro_p"]],
        "Recall": [metrics["macro_r"]],
        "F1-Score": [metrics["macro_f1"]],
        "Support": [sum(metrics["support"])]
    })

    # Add accuracy
    accuracy_row = pd.DataFrame({
        "Category": ["Accuracy"],
        "Precision": [metrics["accuracy"]],
        "Recall": [metrics["accuracy"]],
        "F1-Score": [metrics["accuracy"]],
        "Support": [sum(metrics["support"])]
    })

    report_df = pd.concat([df, macro_row, accuracy_row], ignore_index=True)
    report_df.to_csv(output_path, index=False)


def main():
    # Ensure output directory exists
    output_dir = "./Ten-category_reports"
    os.makedirs(output_dir, exist_ok=True)

    # Load data
    input_file = "Ten-category_result.json"
    data = load_data(input_file)

    # Calculate metrics
    metrics = calculate_metrics(data)

    # Generate standard report
    metrics_path = os.path.join(output_dir, "classification_metrics.csv")
    cm_path = os.path.join(output_dir, "confusion_matrix.png")

    save_metrics_report(metrics, metrics_path)
    generate_confusion_matrix(metrics["true_labels"], metrics["pred_labels"], cm_path)

    print(f"Report generation complete, saved to directory: {output_dir}")
    print(f"Metrics file: classification_metrics.csv")
    print(f"Confusion matrix: confusion_matrix.png")


if __name__ == "__main__":
    main()