evaluate_prototypical.py

"""
Evaluate Prototypical Networks (ResNet50, Swin Transformer) for bearing fault diagnosis.

This script replicates the original evaluation approach using:
- easyfsl library for Prototypical Networks
- Pre-extracted embeddings for efficiency
- ResNet50 or Swin Transformer as feature extractors

Results are saved in format compatible with convert_prototypical_results.py.
"""

import argparse
import os
import pickle
import sys
from pathlib import Path

import numpy as np
import pandas as pd
import torch
import torchvision.models as models
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from torchvision.models import Swin_V2_T_Weights
from PIL import Image
from sklearn.metrics import accuracy_score, precision_recall_fscore_support
import yaml

# Import easyfsl
try:
    from easyfsl.samplers import TaskSampler
    from easyfsl.datasets import FeaturesDataset
    from easyfsl.methods import PrototypicalNetworks
except ImportError:
    print("ERROR: easyfsl not installed. This is required for Prototypical Networks.")
    print("Install with: pip install easyfsl")
    sys.exit(1)


class CWTDataset(Dataset):
    """Custom dataset for CWT bearing fault images."""
    
    def __init__(self, folder_path, transform=None, selected_speed=None, fr_fa_pairs=None):
        """
        Args:
            folder_path: Root directory containing class subdirectories
            transform: Transforms to apply to images
            selected_speed: Filter by specific RPM (None for all speeds)
            fr_fa_pairs: List of (fr, fa) tuples to filter by load conditions
        """
        self.folder_path = folder_path
        self.transform = transform
        self.selected_speed = selected_speed
        self.fr_fa_pairs = fr_fa_pairs or [(None, None)]
        self.images = []
        self.labels = []
        self.label_dict = {'H': 0, 'B': 1, 'IR': 2, 'OR': 3}

        for filename in os.listdir(folder_path):
            if filename.endswith('.png'):
                label = filename.split('_')[0]
                if label in self.label_dict and self.is_valid_image(filename):
                    self.images.append(os.path.join(folder_path, filename))
                    self.labels.append(self.label_dict[label])
        
        print(f"Loaded {len(self.images)} images")

    def is_valid_image(self, filename):
        """Check if image matches filtering criteria."""
        condition, rpm, fr, fa = self.extract_metadata_from_filename(filename)
        for fr_target, fa_target in self.fr_fa_pairs:
            if self.selected_speed is None or rpm == self.selected_speed:
                if (fr == str(fr_target) and fa == str(fa_target)):
                    return True
        return False

    def extract_metadata_from_filename(self, filename):
        """Extract metadata (condition, rpm, fr, fa) from filename."""
        parts = filename.split('_')
        if len(parts) >= 5:
            condition = parts[0]  # 'B', 'H', 'IR', 'OR'
            rpm = parts[1].replace('rpm', '')
            fr = parts[2].replace('kN', '')
            fa = parts[3].replace('kN', '')
            return condition, rpm, fr, fa
        return None, None, None, None

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        image_path = self.images[idx]
        label = self.labels[idx]
        image = Image.open(image_path).convert('RGB')
        if self.transform:
            image = self.transform(image)
        return image, label


def extract_embeddings(dataset, model, device, batch_size=32):
    """
    Extract feature embeddings using the backbone model.
    
    Args:
        dataset: Dataset to extract embeddings from
        model: Feature extractor model
        device: Device to run on
        batch_size: Batch size for extraction
    
    Returns:
        embeddings, labels tensors
    """
    loader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
    embeddings = []
    labels = []
    
    model.eval()
    print(f"Extracting embeddings from {len(dataset)} images...")
    
    with torch.no_grad():
        for inputs, batch_labels in loader:
            inputs = inputs.to(device)
            features = model(inputs)
            embeddings.append(features.cpu())
            labels.append(batch_labels.cpu())
    
    embeddings = torch.cat(embeddings, 0)
    labels = torch.cat(labels, 0)
    
    print(f"✓ Extracted embeddings: {embeddings.shape}")
    return embeddings, labels


def save_features_dataset(dataset, file_path):
    """Save FeaturesDataset to disk."""
    with open(file_path, "wb") as f:
        pickle.dump(dataset, f)
    print(f"✓ Saved features to: {file_path}")


def load_features_dataset(file_path):
    """Load FeaturesDataset from disk."""
    with open(file_path, "rb") as f:
        dataset = pickle.load(f)
    print(f"✓ Loaded features from: {file_path}")
    return dataset


def evaluate_model(model, loader):
    """
    Evaluate Prototypical Network on episodic tasks.
    
    Args:
        model: Prototypical Network model
        loader: DataLoader with TaskSampler
    
    Returns:
        List of [accuracy, precision, recall, f1] for each episode
    """
    results_per_repetition = []
    
    for batch in loader:
        support_images, support_labels, query_images, query_labels, _ = batch
        
        # Process support set to compute prototypes
        model.process_support_set(support_images.to(device), support_labels.to(device))
        
        # Predict on query set
        predictions = model(query_images.to(device))
        
        # Compute metrics
        true_labels = query_labels.cpu().numpy()
        predicted_labels = predictions.argmax(dim=1).cpu().numpy()
        
        acc = accuracy_score(true_labels, predicted_labels)
        prec, rec, f1, _ = precision_recall_fscore_support(
            true_labels, predicted_labels, 
            average='macro', 
            zero_division=0
        )
        
        results_per_repetition.append([acc, prec, rec, f1])
    
    return results_per_repetition


def evaluate_prototypical(
    model_name: str,
    n_shot: int,
    n_tasks: int = 10,
    n_query: int = 1,
    data_dir: str = "Matlab/cwt_images_matlab",
    output_dir: str = "results",
    selected_speed: str = None,
    fr_values: list = None,
    fa_values: list = None,
    force_recompute: bool = False
):
    """
    Evaluate Prototypical Network following the original approach.
    
    Args:
        model_name: Backbone model ('resnet50' or 'swin_v2_t')
        n_shot: Number of support samples per class
        n_tasks: Number of evaluation tasks/episodes
        n_query: Number of query samples per task
        data_dir: Directory containing images
        output_dir: Directory to save results
        selected_speed: Filter by specific RPM (None for all speeds)
        fr_values: List of FR load values
        fa_values: List of FA load values
        force_recompute: Force recomputation of embeddings
    """
    global device
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Using device: {device}")
    
    # Default load conditions
    if fr_values is None:
        fr_values = [124.8]
    if fa_values is None:
        fa_values = [0]
    fr_fa_pairs = list(zip(fr_values, fa_values))
    
    # Data transforms
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                           std=[0.229, 0.224, 0.225])
    ])
    
    # Create dataset
    print(f"\nLoading dataset from: {data_dir}")
    print(f"Speed filter: {selected_speed if selected_speed else 'all_speeds'}")
    print(f"Load conditions (FR, FA): {fr_fa_pairs}")
    
    dataset = CWTDataset(
        folder_path=data_dir,
        transform=transform,
        selected_speed=selected_speed,
        fr_fa_pairs=fr_fa_pairs
    )
    
    if len(dataset) == 0:
        raise ValueError(f"No images found matching criteria")
    
    # Load or create backbone model
    print(f"\nInitializing {model_name} backbone...")
    
    if model_name == 'resnet50':
        model = models.resnet50(pretrained=True)
        model.fc = torch.nn.Flatten()
        model_display = 'Resnet50'
    elif model_name == 'swin_v2_t':
        model = models.swin_v2_t(weights=Swin_V2_T_Weights.IMAGENET1K_V1)
        model.head = torch.nn.Flatten()
        model_display = 'SwinTransformer'
    else:
        raise ValueError(f"Unknown model: {model_name}")
    
    model = model.to(device)
    
    # Define class names
    class_names = ['H', 'B', 'IR', 'OR']
    n_way = len(class_names)
    
    # Features dataset path
    features_dataset_path = f"features_dataset_{model_display}.pkl"
    
    # Extract or load embeddings
    if not os.path.exists(features_dataset_path) or force_recompute:
        print("\nExtracting embeddings...")
        embeddings, labels = extract_embeddings(dataset, model, device)
        features_dataset = FeaturesDataset(labels.tolist(), embeddings, class_names)
        save_features_dataset(features_dataset, features_dataset_path)
    else:
        features_dataset = load_features_dataset(features_dataset_path)
    
    # Initialize Prototypical Network classifier
    few_shot_classifier = PrototypicalNetworks(backbone=torch.nn.Identity())
    
    # Create task sampler
    print(f"\nEvaluating {n_way}-way {n_shot}-shot with {n_tasks} tasks...")
    
    task_sampler = TaskSampler(
        features_dataset,
        n_way=n_way,
        n_shot=n_shot,
        n_query=n_query,
        n_tasks=n_tasks
    )
    
    features_loader = DataLoader(
        features_dataset,
        batch_sampler=task_sampler,
        collate_fn=task_sampler.episodic_collate_fn
    )
    
    # Evaluate
    results_per_repetition = evaluate_model(few_shot_classifier, features_loader)
    
    # Convert to DataFrame
    df_results = pd.DataFrame(
        results_per_repetition, 
        columns=['Accuracy', 'Precision', 'Recall', 'F1']
    )
    
    # Save results
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    
    output_file = output_path / f'Prototypical_{model_display}_{n_shot}-shot_classification_metrics.xlsx'
    
    with pd.ExcelWriter(output_file, engine='openpyxl') as writer:
        # All results
        df_results.to_excel(writer, sheet_name='All_Results', index=False)
        
        # Summary statistics
        summary = df_results.agg({
            'Accuracy': ['mean', 'std'],
            'Precision': ['mean', 'std'],
            'Recall': ['mean', 'std'],
            'F1': ['mean', 'std']
        })
        summary.to_excel(writer, sheet_name='Summary')
    
    print(f"\n✓ Results saved to: {output_file}")
    print("\nSummary Statistics:")
    print(df_results.describe())
    
    return df_results


def main():
    parser = argparse.ArgumentParser(
        description="Evaluate Prototypical Networks for bearing fault diagnosis"
    )
    parser.add_argument(
        '--model',
        type=str,
        required=True,
        choices=['resnet50', 'swin_v2_t'],
        help='Backbone model architecture'
    )
    parser.add_argument(
        '--n-shots',
        type=int,
        nargs='+',
        default=[1, 5, 10],
        help='List of n-shot values to evaluate'
    )
    parser.add_argument(
        '--n-tasks',
        type=int,
        default=10,
        help='Number of evaluation tasks per configuration (default: 10 for t-Student CI)'
    )
    parser.add_argument(
        '--n-query',
        type=int,
        default=1,
        help='Number of query samples per task'
    )
    parser.add_argument(
        '--data-dir',
        type=str,
        default='Matlab/cwt_images_matlab',
        help='Directory containing image data'
    )
    parser.add_argument(
        '--output-dir',
        type=str,
        default='results',
        help='Directory to save results'
    )
    parser.add_argument(
        '--speed',
        type=str,
        default=None,
        help='Filter by specific RPM (e.g., "607")'
    )
    parser.add_argument(
        '--force-recompute',
        action='store_true',
        help='Force recomputation of embeddings'
    )
    
    args = parser.parse_args()
    
    print("=" * 80)
    print("Prototypical Network Evaluation")
    print("=" * 80)
    print(f"\nModel: {args.model}")
    print(f"N-shot values: {args.n_shots}")
    print(f"Tasks per config: {args.n_tasks}")
    print(f"Query samples: {args.n_query}")
    print(f"Data directory: {args.data_dir}")
    print(f"Output directory: {args.output_dir}")
    print()
    
    # Evaluate for each n-shot value
    for n_shot in args.n_shots:
        print(f"\n{'=' * 80}")
        print(f"Evaluating {n_shot}-shot...")
        print('=' * 80)
        
        evaluate_prototypical(
            model_name=args.model,
            n_shot=n_shot,
            n_tasks=args.n_tasks,
            n_query=args.n_query,
            data_dir=args.data_dir,
            output_dir=args.output_dir,
            selected_speed=args.speed,
            force_recompute=args.force_recompute
        )
    
    print(f"\n{'=' * 80}")
    print("Evaluation Complete!")
    print('=' * 80)
    print(f"\nNext steps:")
    print(f"1. Convert results: python convert_prototypical_results.py")
    print(f"2. Aggregate all results: python aggregate_results.py")
    print(f"3. Generate plots: python plot_results.py")


if __name__ == "__main__":
    main()