analyze_true_clustering.py

#!/usr/bin/env python
"""
True Label Clustering Analysis Script

Analyzes clustering metrics (silhouette score, purity, NMI) when using ground truth labels as clusters.
This provides a baseline for perfect clustering performance.

Reads experiment config from results folder and saves analysis results there.

Usage:
    python analyze_true_clustering.py --config results/MVTec/HierarchicalPatchCore/greedy0.1-layer4/all/config.yaml
    python analyze_true_clustering.py --config results/MPDD/HierarchicalPatchCore/greedy0.1-layer4/all/config.yaml
"""

import argparse
import os
import json
import numpy as np
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tqdm import tqdm
from sklearn.metrics import silhouette_score, silhouette_samples, normalized_mutual_info_score
import yaml

import backbones
from dataset import create_dataset, DatasetSplit


def load_config(config_path):
    """Load experiment config from yaml file."""
    with open(config_path, 'r') as f:
        config = yaml.safe_load(f)
    return config


def extract_semantic_features(dataloader, backbone, device, semantic_layer='layer4'):
    """Extract semantic features using specified backbone layer."""
    features_buffer = {}
    
    def hook_fn(module, input, output):
        features_buffer['feat'] = output
    
    layer = getattr(backbone, semantic_layer)
    handle = layer.register_forward_hook(hook_fn)
    
    all_features = []
    all_classes = []
    
    with torch.no_grad():
        for batch in tqdm(dataloader, desc='Extracting features'):
            images = batch['image'].to(device)
            classes = batch['classname']
            
            _ = backbone(images)
            
            feat = features_buffer['feat']
            feat = F.adaptive_avg_pool2d(feat, (1, 1))
            feat = feat.view(feat.size(0), -1)
            
            all_features.append(feat.cpu().numpy())
            all_classes.extend(classes)
    
    handle.remove()
    return np.vstack(all_features), all_classes


def compute_purity(clusters, labels, per_sample_sil=None):
    """Compute clustering purity.

    If per_sample_sil (array of silhouette values per sample) is provided, compute
    the mean silhouette per cluster and include it in the returned cluster_info.
    """
    unique_labels = sorted(set(labels))
    label_to_int = {l: i for i, l in enumerate(unique_labels)}
    labels_int = np.array([label_to_int[l] for l in labels])

    total_correct = 0
    cluster_info = []

    for cluster_id in np.unique(clusters):
        mask = clusters == cluster_id
        cluster_labels = labels_int[mask]

        if len(cluster_labels) == 0:
            continue

        counts = np.bincount(cluster_labels, minlength=len(unique_labels))
        dominant_count = counts.max()
        dominant_label = unique_labels[counts.argmax()]
        total_correct += dominant_count

        info = {
            'cluster_id': int(cluster_id),
            'size': int(mask.sum()),
            'dominant_class': dominant_label,
            'purity': float(dominant_count / mask.sum() * 100)
        }

        if per_sample_sil is not None:
            # average silhouette of samples in this cluster
            sil_vals = per_sample_sil[mask]
            # If cluster has only one sample, silhouette is not defined; set to NaN
            if len(sil_vals) == 0:
                info['silhouette'] = None
            else:
                info['silhouette'] = float(np.nanmean(sil_vals))
        else:
            info['silhouette'] = None

        cluster_info.append(info)

    overall_purity = total_correct / len(labels) * 100
    return overall_purity, cluster_info


def compute_nmi(clusters, labels):
    """Compute Normalized Mutual Information."""
    unique_labels = sorted(set(labels))
    label_to_int = {l: i for i, l in enumerate(unique_labels)}
    labels_int = np.array([label_to_int[l] for l in labels])
    return normalized_mutual_info_score(labels_int, clusters)


def analyze_true_clustering(features, labels):
    """Analyze clustering metrics using ground truth labels as clusters."""
    print('\nAnalyzing true label clustering (baseline)...')
    
    # Convert string labels to integer cluster IDs
    unique_labels = sorted(set(labels))
    label_to_cluster = {label: i for i, label in enumerate(unique_labels)}
    clusters = np.array([label_to_cluster[label] for label in labels])
    
    # Compute metrics
    try:
        silhouette = silhouette_score(features, clusters)
    except Exception as e:
        print(f"Warning: Could not compute silhouette score: {e}")
        silhouette = None
    
    purity, cluster_info = compute_purity(clusters, labels)
    nmi = compute_nmi(clusters, labels)
    
    print('\n' + '=' * 80)
    print('True Label Clustering Results (Baseline)')
    print('=' * 80)
    print(f'Silhouette Score: {silhouette:.4f}' if silhouette is not None else 'Silhouette Score: N/A')
    print(f'Purity:           {purity:.2f}%')
    print(f'NMI:              {nmi:.4f}')
    print('=' * 80)
    
    # Per-sample silhouette for cluster info
    if silhouette is not None:
        try:
            per_sample_sil = silhouette_samples(features, clusters)
            _, cluster_info = compute_purity(clusters, labels, per_sample_sil)
        except Exception:
            per_sample_sil = None
    
    print(f'\nCluster Distribution (Ground Truth Classes):')
    print('-' * 60)
    for info in sorted(cluster_info, key=lambda x: x['cluster_id']):
        sil_str = f", silhouette: {info['silhouette']:.3f}" if info['silhouette'] is not None else ""
        print(f"  Cluster {info['cluster_id']:3d}: {info['size']:4d} samples, "
              f"dominant: {info['dominant_class']:15s}, purity: {info['purity']:.1f}%{sil_str}")
    
    return {
        'silhouette': silhouette,
        'purity': purity,
        'nmi': nmi
    }, clusters, cluster_info


def save_results(save_path, results, cluster_info, labels):
    """Save analysis results to JSON file."""
    unique_labels = sorted(set(labels))
    
    output = {
        'num_samples': len(labels),
        'num_ground_truth_classes': len(unique_labels),
        'ground_truth_classes': unique_labels,
        'results': {
            'silhouette': float(results['silhouette']) if results['silhouette'] is not None else None,
            'purity': float(results['purity']),
            'nmi': float(results['nmi'])
        },
        'cluster_distribution': [
            {
                'cluster_id': int(info['cluster_id']),
                'size': int(info['size']),
                'dominant_class': info['dominant_class'],
                'purity': float(info['purity']),
                'silhouette': (float(info['silhouette']) if ('silhouette' in info and info['silhouette'] is not None) else None)
            }
            for info in sorted(cluster_info, key=lambda x: x['cluster_id'])
        ]
    }
    
    with open(save_path, 'w') as f:
        json.dump(output, f, indent=2)
    
    return output


def main():
    parser = argparse.ArgumentParser(description='Analyze true label clustering (baseline) on semantic features')
    parser.add_argument('--config', type=str, required=True,
                        help='Path to experiment config.yaml file')
    parser.add_argument('--device', type=str, default='cuda',
                        help='Device (cuda or cpu)')
    args = parser.parse_args()
    
    # Load config
    config = load_config(args.config)
    config_dir = os.path.dirname(args.config)
    
    # Extract parameters from config
    dataset_name = config['DATASET']['name']
    data_path = config['DATASET']['datadir']
    resize = config['DATASET']['resize']
    imagesize = config['DATASET']['imagesize']
    classname = config['DATASET']['classname']
    backbone_name = config['MODEL']['backbone']
    semantic_layer = config['MODEL'].get('semantic_layer', 'layer4')
    batch_size = config['TRAIN'].get('test_batch_size', 64)
    num_workers = config['TRAIN'].get('num_workers', 8)
    
    device = torch.device(args.device if torch.cuda.is_available() else 'cpu')
    
    print('=' * 80)
    print('True Label Clustering Analysis (Baseline)')
    print('=' * 80)
    print(f'Config:         {args.config}')
    print(f'Dataset:        {dataset_name}')
    print(f'Data path:      {data_path}')
    print(f'Classname:      {classname}')
    print(f'Resize:         {resize}')
    print(f'Image size:     {imagesize}')
    print(f'Backbone:       {backbone_name}')
    print(f'Semantic layer: {semantic_layer}')
    print(f'Device:         {device}')
    print()
    
    # Load dataset
    print('Loading dataset...')
    dataset = create_dataset(
        dataname=dataset_name,
        source=data_path,
        classname=classname,
        resize=resize,
        imagesize=imagesize,
        split=DatasetSplit.TRAIN
    )
    
    dataloader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=True
    )
    
    print(f'Total samples: {len(dataset)}')
    
    # Load backbone
    print(f'\nLoading backbone ({backbone_name})...')
    backbone = backbones.load(backbone_name)
    backbone = backbone.to(device)
    backbone.eval()
    
    # Extract features
    print(f'\nExtracting semantic features from {semantic_layer}...')
    features, labels = extract_semantic_features(
        dataloader, backbone, device, semantic_layer
    )
    print(f'Features shape: {features.shape}')
    print(f'Unique classes: {len(set(labels))}')
    
    # Analyze true label clustering
    results, clusters, cluster_info = analyze_true_clustering(features, labels)
    
    # Save results
    save_path = os.path.join(config_dir, 'true_label_analysis.json')
    output = save_results(save_path, results, cluster_info, labels)
    
    print('\n' + '=' * 80)
    print(f'Results saved to: {save_path}')
    print('=' * 80)


if __name__ == '__main__':
    main()