analyze_kmeans_clustering.py

#!/usr/bin/env python
"""
K-means Clustering Analysis Script

Analyzes K-means clustering on semantic features extracted from anomaly detection datasets.
Computes silhouette score, purity, and NMI for different numbers of clusters (k).

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

Usage:
    python analyze_kmeans_clustering.py --config results/MVTec/HierarchicalPatchCore/greedy0.1-layer4/all/config.yaml
    python analyze_kmeans_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
from sklearn.cluster import KMeans
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_kmeans(features, labels, k_range=range(2, 16)):
    """Run K-means clustering for different k values and analyze results."""
    print('\nRunning K-means clustering...')
    
    results = []
    best_silhouette = -1
    best_k = -1
    
    print('\n' + '=' * 80)
    print(f'{"k":^10} | {"Silhouette":^12} | {"Purity":^10} | {"NMI":^10}')
    print('-' * 80)
    
    for k in k_range:
        kmeans = KMeans(n_clusters=k, random_state=42, n_init=10)
        clusters = kmeans.fit_predict(features)
        
        silhouette = silhouette_score(features, clusters)
        purity, _ = compute_purity(clusters, labels)
        nmi = compute_nmi(clusters, labels)
        
        results.append({
            'k': k,
            'silhouette': silhouette,
            'purity': purity,
            'nmi': nmi
        })
        
        marker = ''
        if silhouette > best_silhouette:
            best_silhouette = silhouette
            best_k = k
            marker = ' <-- BEST'
        
        print(f'{k:^10} | {silhouette:^12.4f} | {purity:^10.2f} | {nmi:^10.4f}{marker}')
    
    print('=' * 80)
    
    # Best k details
    best_result = next(r for r in results if r['k'] == best_k)
    print(f'\nBest k (max silhouette): {best_k}')
    print(f'  - Silhouette Score:   {best_result["silhouette"]:.4f}')
    print(f'  - Purity:             {best_result["purity"]:.2f}%')
    print(f'  - NMI:                {best_result["nmi"]:.4f}')
    
    # Cluster distribution for best k (include per-sample silhouette -> per-cluster mean)
    kmeans = KMeans(n_clusters=best_k, random_state=42, n_init=10)
    best_clusters = kmeans.fit_predict(features)
    
    # compute per-sample silhouette for the best k
    try:
        per_sample_sil = silhouette_samples(features, best_clusters)
    except Exception:
        # fallback: if silhouette_samples fails (e.g., single cluster), set to None
        per_sample_sil = None

    _, cluster_info = compute_purity(best_clusters, labels, per_sample_sil)
    
    print(f'\nCluster Distribution (k={best_k}):')
    print('-' * 60)
    for info in sorted(cluster_info, key=lambda x: x['cluster_id']):
        print(f"  Cluster {info['cluster_id']:3d}: {info['size']:4d} samples, "
              f"dominant: {info['dominant_class']:15s}, purity: {info['purity']:.1f}%")
    
    return results, best_clusters, cluster_info


def save_results(save_path, results, cluster_info, labels):
    """Save analysis results to JSON file."""
    unique_labels = sorted(set(labels))
    
    # Find best k (max silhouette)
    best_idx = max(range(len(results)), key=lambda i: results[i]['silhouette'])
    best_result = results[best_idx]
    
    output = {
        'num_samples': len(labels),
        'num_ground_truth_classes': len(unique_labels),
        'ground_truth_classes': unique_labels,
        'k_values': [r['k'] for r in results],
        'results': [
            {
                'k': r['k'],
                'silhouette': float(r['silhouette']),
                'purity': float(r['purity']),
                'nmi': float(r['nmi'])
            }
            for r in results
        ],
        'best_k': {
            'k': int(best_result['k']),
            'silhouette': float(best_result['silhouette']),
            'purity': float(best_result['purity']),
            'nmi': float(best_result['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 K-means clustering 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)')
    parser.add_argument('--k_min', type=int, default=2,
                        help='Minimum number of clusters')
    parser.add_argument('--k_max', type=int, default=15,
                        help='Maximum number of clusters')
    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('K-means Clustering Analysis')
    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(f'K range:        {args.k_min} to {args.k_max}')
    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 K-means clustering
    k_range = range(args.k_min, args.k_max + 1)
    results, best_clusters, cluster_info = analyze_kmeans(features, labels, k_range)
    
    # Save results
    save_path = os.path.join(config_dir, 'kmeans_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()