calculate_metrics.py

#!/usr/bin/env python3
"""
Calculate evaluation metrics for VQA predictions.
Metrics: BLEU@2, ROUGE-L, METEOR, CIDEr, SPICE, BERTScore, VQA Accuracy
"""

import json
import numpy as np
from collections import Counter
import re
import string
from typing import List, Dict, Tuple, Any
import pandas as pd

def install_required_packages():
    """Install required packages if not available"""
    import subprocess
    import sys
    
    packages = [
        'nltk',
        'rouge-score',
        'pycocoevalcap',
        'bert-score',
        'transformers',
        'torch'
    ]
    
    for package in packages:
        try:
            __import__(package.replace('-', '_'))
        except ImportError:
            print(f"Installing {package}...")
            subprocess.check_call([sys.executable, "-m", "pip", "install", package])

# Install packages first
try:
    install_required_packages()
except Exception as e:
    print(f"Error installing packages: {e}")
    print("Please install manually: pip install nltk rouge-score pycocoevalcap bert-score transformers torch")

import nltk
from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
from rouge_score import rouge_scorer
from bert_score import score as bert_score

# Download required NLTK data
try:
    nltk.download('punkt', quiet=True)
    nltk.download('wordnet', quiet=True)
    nltk.download('omw-1.4', quiet=True)
except:
    pass

def normalize_text(text: str) -> str:
    """Normalize Vietnamese text for evaluation"""
    if not isinstance(text, str):
        text = str(text)
    
    # Remove U+202F (narrow no-break space) character that causes tokenization issues
    text = text.replace('\u202F', ' ')
    
    # Convert to lowercase
    text = text.lower()
    
    # Remove punctuation
    text = text.translate(str.maketrans('', '', string.punctuation))
    
    # Remove extra whitespace
    text = ' '.join(text.split())
    
    return text

def tokenize_vietnamese(text: str) -> List[str]:
    """Simple tokenization for Vietnamese text"""
    text = normalize_text(text)
    return text.split()

def calculate_bleu1(reference: str, candidate: str) -> float:
    """Calculate BLEU-1 score"""
    ref_tokens = tokenize_vietnamese(reference)
    cand_tokens = tokenize_vietnamese(candidate)
    
    if len(cand_tokens) == 0:
        return 0.0
    
    smoothing = SmoothingFunction()
    weights = (1.0, 0, 0, 0)  # BLEU-1 weights
    
    try:
        score = sentence_bleu([ref_tokens], cand_tokens, weights=weights, 
                            smoothing_function=smoothing.method1)
        return float(score) if isinstance(score, (int, float)) else 0.0
    except:
        return 0.0

def calculate_bleu2(reference: str, candidate: str) -> float:
    """Calculate BLEU-2 score"""
    ref_tokens = tokenize_vietnamese(reference)
    cand_tokens = tokenize_vietnamese(candidate)
    
    if len(cand_tokens) == 0:
        return 0.0
    
    smoothing = SmoothingFunction()
    weights = (0.5, 0.5, 0, 0)  # BLEU-2 weights
    
    try:
        score = sentence_bleu([ref_tokens], cand_tokens, weights=weights, 
                            smoothing_function=smoothing.method1)
        return float(score) if isinstance(score, (int, float)) else 0.0
    except:
        return 0.0

def calculate_bleu3(reference: str, candidate: str) -> float:
    """Calculate BLEU-3 score"""
    ref_tokens = tokenize_vietnamese(reference)
    cand_tokens = tokenize_vietnamese(candidate)
    
    if len(cand_tokens) == 0:
        return 0.0
    
    smoothing = SmoothingFunction()
    weights = (1/3, 1/3, 1/3, 0)  # BLEU-3 weights
    
    try:
        score = sentence_bleu([ref_tokens], cand_tokens, weights=weights, 
                            smoothing_function=smoothing.method1)
        return float(score) if isinstance(score, (int, float)) else 0.0
    except:
        return 0.0

def calculate_bleu4(reference: str, candidate: str) -> float:
    """Calculate BLEU-4 score"""
    ref_tokens = tokenize_vietnamese(reference)
    cand_tokens = tokenize_vietnamese(candidate)
    
    if len(cand_tokens) == 0:
        return 0.0
    
    smoothing = SmoothingFunction()
    weights = (0.25, 0.25, 0.25, 0.25)  # BLEU-4 weights
    
    try:
        score = sentence_bleu([ref_tokens], cand_tokens, weights=weights, 
                            smoothing_function=smoothing.method1)
        return float(score) if isinstance(score, (int, float)) else 0.0
    except:
        return 0.0

def calculate_rouge_l(reference: str, candidate: str) -> float:
    """Calculate ROUGE-L score"""
    scorer = rouge_scorer.RougeScorer(['rougeL'], use_stemmer=False)
    scores = scorer.score(reference, candidate)
    return scores['rougeL'].fmeasure

def calculate_meteor(reference: str, candidate: str) -> float:
    """Calculate METEOR score"""
    try:
        from nltk.translate.meteor_score import meteor_score
        ref_tokens = tokenize_vietnamese(reference)
        cand_tokens = tokenize_vietnamese(candidate)
        
        if len(cand_tokens) == 0 or len(ref_tokens) == 0:
            return 0.0
            
        return meteor_score([ref_tokens], cand_tokens)
    except:
        # Fallback to word overlap if METEOR fails
        ref_tokens = set(tokenize_vietnamese(reference))
        cand_tokens = set(tokenize_vietnamese(candidate))
        
        if len(ref_tokens) == 0:
            return 0.0
            
        overlap = len(ref_tokens.intersection(cand_tokens))
        return overlap / len(ref_tokens)

def calculate_cider_proper(references: List[str], candidates: List[str]) -> float:
    """Calculate CIDEr score using pycocoevalcap"""
    try:
        from pycocoevalcap.cider.cider import Cider
        
        # Format data for pycocoevalcap
        gts = {}  # ground truth
        res = {}  # results
        
        for i, (ref, cand) in enumerate(zip(references, candidates)):
            gts[i] = [ref]  # pycocoevalcap expects list of references
            res[i] = [cand]
        
        # Calculate CIDEr
        cider_scorer = Cider()
        score, scores = cider_scorer.compute_score(gts, res)
        
        return float(score)
    except Exception as e:
        print(f"CIDEr calculation failed: {e}")
        print("Falling back to simplified CIDEr calculation...")
        return calculate_cider_simplified_fallback(references, candidates)

def calculate_cider_simplified_fallback(references: List[str], candidates: List[str]) -> float:
    """Fallback simplified CIDEr score calculation"""
    scores = []
    
    for ref, cand in zip(references, candidates):
        ref_tokens = tokenize_vietnamese(ref)
        cand_tokens = tokenize_vietnamese(cand)
        
        if len(cand_tokens) == 0:
            scores.append(0.0)
            continue
            
        # Calculate n-gram overlap (up to 4-grams)
        ngram_scores = []
        for n in range(1, 5):
            ref_ngrams = [' '.join(ref_tokens[i:i+n]) for i in range(len(ref_tokens)-n+1)]
            cand_ngrams = [' '.join(cand_tokens[i:i+n]) for i in range(len(cand_tokens)-n+1)]
            
            if len(ref_ngrams) == 0 or len(cand_ngrams) == 0:
                ngram_scores.append(0.0)
                continue
                
            ref_counter = Counter(ref_ngrams)
            cand_counter = Counter(cand_ngrams)
            
            overlap = sum((ref_counter & cand_counter).values())
            total = sum(cand_counter.values())
            
            if total > 0:
                ngram_scores.append(overlap / total)
            else:
                ngram_scores.append(0.0)
        
        scores.append(np.mean(ngram_scores))
    
    return float(np.mean(scores))

def calculate_spice_proper(references: List[str], candidates: List[str]) -> float:
    """Calculate SPICE score using pycocoevalcap"""
    try:
        from pycocoevalcap.spice.spice import Spice
        
        # Format data for pycocoevalcap
        gts = {}  # ground truth
        res = {}  # results        
        
        for i, (ref, cand) in enumerate(zip(references, candidates)):
            if len(cand.split()) > 64:
                res[i] = "không rõ"
            else:
                res[i] = [cand]

            gts[i] = [ref]  # pycocoevalcap expects list of references
            res[i] = [cand]
        
        # Calculate SPICE
        spice_scorer = Spice()
        score, scores = spice_scorer.compute_score(gts, res)
        
        return float(score)
    except Exception as e:
        print(f"SPICE calculation failed: {e}")
        print("Falling back to simplified SPICE calculation...")
        return calculate_spice_simplified_fallback(references, candidates)

def calculate_spice_simplified_fallback(references: List[str], candidates: List[str]) -> float:
    """Fallback simplified SPICE score (word overlap based)"""
    scores = []
    for ref, cand in zip(references, candidates):
        ref_tokens = set(tokenize_vietnamese(ref))
        cand_tokens = set(tokenize_vietnamese(cand))
        
        if len(ref_tokens) == 0 and len(cand_tokens) == 0:
            scores.append(1.0)
        elif len(ref_tokens) == 0 or len(cand_tokens) == 0:
            scores.append(0.0)
        else:
            overlap = len(ref_tokens.intersection(cand_tokens))
            union = len(ref_tokens.union(cand_tokens))
            scores.append(overlap / union if union > 0 else 0.0)
    
    return float(np.mean(scores))

def calculate_bert_score(references: List[str], candidates: List[str]) -> float:
    """Calculate BERTScore F1"""
    try:
        # Use multilingual BERT for Vietnamese
        P, R, F1 = bert_score(candidates, references, lang='vi', verbose=False)
        return F1.mean().item()
    except Exception as e:
        print(f"BERTScore calculation failed: {e}")
        # Fallback to word overlap
        scores = []
        for ref, cand in zip(references, candidates):
            ref_tokens = set(tokenize_vietnamese(ref))
            cand_tokens = set(tokenize_vietnamese(cand))
            
            if len(ref_tokens) == 0 and len(cand_tokens) == 0:
                scores.append(1.0)
            elif len(ref_tokens) == 0 or len(cand_tokens) == 0:
                scores.append(0.0)
            else:
                overlap = len(ref_tokens.intersection(cand_tokens))
                total = max(len(ref_tokens), len(cand_tokens))
                scores.append(overlap / total)
        
        return float(np.mean(scores))

def calculate_vqa_accuracy(references: List[str], candidates: List[str]) -> float:
    """Calculate VQA accuracy (exact match after normalization)"""
    correct = 0
    total = len(references)
    
    for ref, cand in zip(references, candidates):
        ref_norm = normalize_text(ref)
        cand_norm = normalize_text(cand)
        
        if ref_norm == cand_norm:
            correct += 1
    
    return correct / total if total > 0 else 0.0

def load_predictions(file_path: str) -> Tuple[List[str], List[str]]:
    """Load predictions from JSON file"""
    with open(file_path, 'r', encoding='utf-8') as f:
        data = json.load(f)
    
    references = []
    candidates = []
    
    for item in data:
        # Remove U+202F character from both reference and prediction
        ref = item['gt'].replace('\u202F', ' ') if isinstance(item['gt'], str) else str(item['gt']).replace('\u202F', ' ')
        pred = item['pr'].replace('\u202F', ' ') if isinstance(item['pr'], str) else str(item['pr']).replace('\u202F', ' ')
        
        references.append(ref)
        candidates.append(pred)
    
    return references, candidates

def main():
    """Main function to calculate all metrics"""
    predictions_file = "/home/vlai-gpt-oss/GPT_DAM_VQA/evaluation_results/eval_random_gemini_20250905_023046/predictions_gemini.json"
    
    print("Loading predictions...")
    references, candidates = load_predictions(predictions_file)
    
    print(f"Loaded {len(references)} prediction pairs")
    
    # Calculate metrics
    print("\nCalculating metrics...")
    
    # BLEU scores (1-4)
    print("Calculating BLEU-1...")
    bleu1_scores = [calculate_bleu1(ref, cand) for ref, cand in zip(references, candidates)]
    bleu1_avg = np.mean(bleu1_scores)
    
    print("Calculating BLEU-2...")
    bleu2_scores = [calculate_bleu2(ref, cand) for ref, cand in zip(references, candidates)]
    bleu2_avg = np.mean(bleu2_scores)
    
    print("Calculating BLEU-3...")
    bleu3_scores = [calculate_bleu3(ref, cand) for ref, cand in zip(references, candidates)]
    bleu3_avg = np.mean(bleu3_scores)
    
    print("Calculating BLEU-4...")
    bleu4_scores = [calculate_bleu4(ref, cand) for ref, cand in zip(references, candidates)]
    bleu4_avg = np.mean(bleu4_scores)
    
    # ROUGE-L
    print("Calculating ROUGE-L...")
    rouge_l_scores = [calculate_rouge_l(ref, cand) for ref, cand in zip(references, candidates)]
    rouge_l_avg = np.mean(rouge_l_scores)
    
    # METEOR
    print("Calculating METEOR...")
    meteor_scores = [calculate_meteor(ref, cand) for ref, cand in zip(references, candidates)]
    meteor_avg = float(np.mean(meteor_scores))
    
    # CIDEr
    print("Calculating CIDEr...")
    cider_avg = float(calculate_cider_proper(references, candidates))
    
    # SPICE
    print("Calculating SPICE...")
    spice_avg = float(calculate_spice_proper(references, candidates))
    
    # BERTScore
    print("Calculating BERTScore...")
    bert_score_avg = calculate_bert_score(references, candidates)
    
    # VQA Accuracy
    print("Calculating VQA Accuracy...")
    vqa_accuracy = calculate_vqa_accuracy(references, candidates)
    
    # Create results table
    results = {
        'Metric': ['BLEU@1', 'BLEU@2', 'BLEU@3', 'BLEU@4', 'ROUGE-L', 'METEOR', 'CIDEr', 'SPICE', 'BERTScore', 'VQA Accuracy'],
        'Score': [
            f"{bleu1_avg:.4f}",
            f"{bleu2_avg:.4f}",
            f"{bleu3_avg:.4f}",
            f"{bleu4_avg:.4f}",
            f"{rouge_l_avg:.4f}",
            f"{meteor_avg:.4f}",
            f"{cider_avg:.4f}",
            f"{spice_avg:.4f}",
            f"{bert_score_avg:.4f}",
            f"{vqa_accuracy:.4f}"
        ]
    }
    
    df = pd.DataFrame(results)
    
    print("\n" + "="*50)
    print("EVALUATION RESULTS")
    print("="*50)
    print(df.to_string(index=False))
    
    # Generate markdown table
    markdown_table = "| Metric | Score |\n|--------|-------|\n"
    for metric, score in zip(results['Metric'], results['Score']):
        markdown_table += f"| {metric} | {score} |\n"
    
    print("\n" + "="*50)
    print("MARKDOWN TABLE")
    print("="*50)
    print(markdown_table)
    
    # Save results to the same directory as the input predictions
    import os
    predictions_dir = os.path.dirname(predictions_file)
    output_file = os.path.join(predictions_dir, "evaluation_metrics.md")
    
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write("# DAM Model Evaluation Results\n\n")
        f.write("## Metrics Summary\n\n")
        f.write(markdown_table)
        f.write(f"\n\n**Total samples evaluated:** {len(references)}\n")
        f.write(f"**Evaluation date:** September 3, 2025\n")
    
    print(f"\nResults saved to: {output_file}")

if __name__ == "__main__":
    main()