analyze_finch_clustering.py

#!/usr/bin/env python
"""
FINCH Clustering Analysis Script

Analyzes FINCH clustering on semantic features extracted from anomaly detection datasets.
Computes silhouette score, purity, and NMI for each partition.

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

Usage:
    python analyze_finch_clustering.py --config results/MVTec/HierarchicalPatchCore/greedy0.1-layer4/all/config.yaml
    python analyze_finch_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 finch import FINCH
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_finch(features, labels):
    """Run FINCH clustering and analyze all partitions."""
    print('\nRunning FINCH clustering...')
    c, num_clust, _ = FINCH(features, distance='cosine')
    num_clust = np.array(num_clust)
    
    print(f'Number of partitions: {len(num_clust)}')
    print(f'Clusters per partition: {num_clust.tolist()}')
    
    results = []
    best_silhouette = -1
    best_partition = -1
    
    print('\n' + '=' * 80)
    print(f'{"Partition":^10} | {"Clusters":^10} | {"Silhouette":^12} | {"Purity":^10} | {"NMI":^10}')
    print('-' * 80)
    
    for i in range(len(num_clust)):
        clusters = c[:, i]
        
        silhouette = silhouette_score(features, clusters)
        purity, _ = compute_purity(clusters, labels)
        nmi = compute_nmi(clusters, labels)
        
        results.append({
            'partition': i,
            'num_clusters': num_clust[i],
            'silhouette': silhouette,
            'purity': purity,
            'nmi': nmi
        })
        
        marker = ''
        if silhouette > best_silhouette:
            best_silhouette = silhouette
            best_partition = i
            marker = ' <-- BEST'
        
        print(f'{i:^10} | {num_clust[i]:^10} | {silhouette:^12.4f} | {purity:^10.2f} | {nmi:^10.4f}{marker}')
    
    print('=' * 80)
    
    # Best partition details
    best_result = results[best_partition]
    print(f'\nBest Partition (max silhouette): {best_partition}')
    print(f'  - Number of Clusters: {best_result["num_clusters"]}')
    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 partition (include per-sample silhouette -> per-cluster mean)
    best_clusters = c[:, best_partition]
    # compute per-sample silhouette for the best partition
    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 (Partition {best_partition}):')
    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, c, num_clust, cluster_info


def save_results(save_path, results, cluster_info, num_clust, 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,
        'num_partitions': len(num_clust),
        'partitions': [
            {
                'partition': r['partition'],
                'num_clusters': int(r['num_clusters']),
                'silhouette': float(r['silhouette']),
                'purity': float(r['purity']),
                'nmi': float(r['nmi'])
            }
            for r in results
        ],
        'best_partition': {
            'partition': int(results[-1]['partition']),  # Last one with highest silhouette
            'num_clusters': int(results[-1]['num_clusters']),
            'silhouette': float(results[-1]['silhouette']),
            'purity': float(results[-1]['purity']),
            'nmi': float(results[-1]['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'])
            ]
        }
    }
    
    # Find best partition (max silhouette)
    best_idx = max(range(len(results)), key=lambda i: results[i]['silhouette'])
    output['best_partition']['partition'] = results[best_idx]['partition']
    output['best_partition']['num_clusters'] = int(results[best_idx]['num_clusters'])
    output['best_partition']['silhouette'] = float(results[best_idx]['silhouette'])
    output['best_partition']['purity'] = float(results[best_idx]['purity'])
    output['best_partition']['nmi'] = float(results[best_idx]['nmi'])
    
    with open(save_path, 'w') as f:
        json.dump(output, f, indent=2)
    
    return output


def main():
    parser = argparse.ArgumentParser(description='Analyze FINCH 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)')
    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('FINCH 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()
    
    # 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 FINCH clustering
    results, clusters, num_clust, cluster_info = analyze_finch(features, labels)
    
    # Save results
    save_path = os.path.join(config_dir, 'finch_analysis.json')
    output = save_results(save_path, results, cluster_info, num_clust, labels)
    
    print('\n' + '=' * 80)
    print(f'Results saved to: {save_path}')
    print('=' * 80)


if __name__ == '__main__':
    main()