analyze_dbscan_clustering.py

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

Analyzes DBSCAN clustering on semantic features extracted from anomaly detection datasets.
Computes silhouette score, purity, and NMI for the clustering result.

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

Usage:
    python analyze_dbscan_clustering.py --config results/MVTec/HierarchicalPatchCore/greedy0.1-layer4-dbscan/all/config.yaml
    python analyze_dbscan_clustering.py --config results/MPDD/HierarchicalPatchCore/greedy0.1-layer4-dbscan/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 DBSCAN
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_dbscan(features, labels, eps=0.5, min_samples=5):
    """Run DBSCAN clustering and analyze results."""
    print('\nRunning DBSCAN clustering...')
    print(f'Parameters: eps={eps}, min_samples={min_samples}')
    
    dbscan = DBSCAN(eps=eps, min_samples=min_samples)
    clusters = dbscan.fit_predict(features)
    
    # Handle noise points (labeled as -1)
    noise_mask = clusters == -1
    num_noise = noise_mask.sum()
    num_clusters = len(set(clusters)) - (1 if -1 in clusters else 0)
    
    print(f'Number of clusters found: {num_clusters}')
    print(f'Number of noise points: {num_noise}')
    
    # If noise points exist, assign them to nearest cluster (similar to original code)
    if num_noise > 0:
        print(f'Assigning {num_noise} noise points to nearest clusters...')
        from sklearn.metrics.pairwise import euclidean_distances
        
        # Get core samples (non-noise points)
        core_indices = np.where(clusters != -1)[0]
        noise_indices = np.where(clusters == -1)[0]
        
        if len(core_indices) > 0:
            # Compute distances from noise points to all points
            distances = euclidean_distances(features[noise_indices], features[core_indices])
            # Find nearest core point for each noise point
            nearest_core_indices = core_indices[np.argmin(distances, axis=1)]
            # Assign cluster of nearest core point
            clusters[noise_indices] = clusters[nearest_core_indices]
        else:
            # If no core samples, assign all to cluster 0
            clusters[noise_indices] = 0
    
    # Recompute after noise assignment
    num_clusters = len(set(clusters))
    print(f'Final number of clusters: {num_clusters}')
    
    # 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('DBSCAN Clustering Results')
    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:')
    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 {
        'eps': eps,
        'min_samples': min_samples,
        'num_clusters': num_clusters,
        'num_noise': num_noise,
        '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,
        'dbscan_params': {
            'eps': float(results['eps']),
            'min_samples': int(results['min_samples'])
        },
        'results': {
            'num_clusters': int(results['num_clusters']),
            'num_noise': int(results['num_noise']),
            '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 DBSCAN 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('--eps', type=float, default=0.5,
                        help='DBSCAN eps parameter')
    parser.add_argument('--min_samples', type=int, default=5,
                        help='DBSCAN min_samples parameter')
    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('DBSCAN 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'DBSCAN eps:     {args.eps}')
    print(f'DBSCAN min_samples: {args.min_samples}')
    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 DBSCAN clustering
    results, clusters, cluster_info = analyze_dbscan(features, labels, eps=args.eps, min_samples=args.min_samples)
    
    # Save results
    save_path = os.path.join(config_dir, 'dbscan_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()