main.py

import os
import json
import pickle
import numpy as np
from tqdm.auto import tqdm
from omegaconf import OmegaConf, ListConfig
from collections import defaultdict

from torch.utils.data import DataLoader

import common
import sampler
import backbones
import utils
from dataset import DatasetSplit, create_dataset
from models import create_model
from metrics import Evaluator, get_binary_results

# load patchcor
def get_sampler(name, percentage, device, seed: int = 0, **params):
    if name == "identity":
        return sampler.IdentitySampler()
    elif name == "greedy_coreset":
        return sampler.GreedyCoresetSampler(
            percentage=percentage, 
            device=device, 
            seed=seed, 
            is_reduce=params.get('reduce', True), 
            batch_size=params.get('batch_size', -1),
            method=params.get('method', 'random_proj'),
            dimension_to_project_features_to=params.get('dimension_to_project_features_to', 128)
        )
    elif name == "approx_greedy_coreset":
        return sampler.ApproximateGreedyCoresetSampler(percentage=percentage, device=device, seed=seed)
    
def calc_metrics(segmentations, masks_gt, class_labels, semantic_match_indices = []):
    evaluator = Evaluator()
    eval_args = {
        'cls_names'              : class_labels,
        'imgs_masks'             : masks_gt,
        'anomaly_maps'           : segmentations,
    }
    
    results = {}
    for cls in tqdm(np.unique(class_labels), desc='Evaluating per class...'):
        results[cls] = evaluator.run(results=eval_args, cls_name=cls)
        
    if len(np.unique(class_labels)) > 1:
        # Always compute global evaluation (ignoring clusters)
        results['all'] = evaluator.run(results=eval_args, cls_name='all')
        
        # Additionally compute per-cluster evaluation if semantic_match_indices exists
        if len(semantic_match_indices) > 0:
            r_all = defaultdict(list)
            skipped_clusters = []
            for idx in np.unique(semantic_match_indices):
                match_idx = np.where(semantic_match_indices==idx)[0]
                try:
                    r = evaluator.run(results=eval_args, cls_name='all', cls_index=match_idx)
                    for k, v in r.items():
                        r_all[k].append(v)
                except AssertionError as e:
                    # Debug info for the failed cluster
                    cluster_masks = masks_gt[match_idx]
                    cluster_segmentations = segmentations[match_idx]
                    n_samples = len(match_idx)
                    n_anomaly = np.sum(cluster_masks.max(axis=(1,2)) > 0)
                    n_normal = n_samples - n_anomaly
                    gt_all_zero = np.all(cluster_masks == 0)
                    anomap_range = (cluster_segmentations.min(), cluster_segmentations.max())
                    
                    print(f"\n[DEBUG] Cluster {idx} failed with AssertionError:")
                    print(f"  - Total samples: {n_samples}")
                    print(f"  - Normal samples (GT all zero): {n_normal}")
                    print(f"  - Anomaly samples: {n_anomaly}")
                    print(f"  - GT mask all zero: {gt_all_zero}")
                    print(f"  - Anomaly map range: {anomap_range}")
                    
                    skipped_clusters.append(idx)
                    continue
            
            if skipped_clusters:
                print(f"Warning: Skipped clusters {skipped_clusters} due to AUPRO calculation error")
            
            if r_all:
                results['all_per_cluster'] = dict([(k, np.mean(v)) for k, v in r_all.items()])
    
    binary_results = get_binary_results(
        eval_args              = eval_args, 
        semantic_match_indices = semantic_match_indices
    )
    
    return results, binary_results, eval_args


def evaluate(model, testloader):
    # prediction
    outputs = model.predict(
        testloader
    )
    
    segmentations = np.array(outputs['masks'])
    masks_gt = np.array(outputs['masks_gt'])
    
    class_labels = np.asarray([
        x[0] for x in testloader.dataset.data_to_iterate
    ])
    
    results, binary_results, eval_args = calc_metrics(
        segmentations = segmentations,
        masks_gt      = masks_gt,
        class_labels  = class_labels,
        semantic_match_indices = outputs.get('semantic_match_indices', [])
    )
    
    if 'semantic_match_indices' in outputs:
        eval_args['semantic_match_indices'] = outputs['semantic_match_indices']
        
    return results, binary_results, eval_args, outputs['infer_time']


def run(cfg):
    utils.fix_seeds(seed=cfg.DEFAULT.seed)
    
    # savedir
    if isinstance(cfg.DATASET.classname, list) or isinstance(cfg.DATASET.classname, ListConfig):
            savedir = os.path.join(cfg.DEFAULT.savedir, '-'.join(cfg.DATASET.classname))
    else:
        savedir = os.path.join(cfg.DEFAULT.savedir, cfg.DATASET.classname)
    os.makedirs(savedir, exist_ok=True)
    
    # device
    device = utils.set_torch_device(gpu_ids=cfg.DEFAULT.device_ids)
    
    trainset = create_dataset(
        dataname  = cfg.DATASET.name,
        source    = cfg.DATASET.datadir,
        classname = cfg.DATASET.classname,
        resize    = cfg.DATASET.resize,
        imagesize = cfg.DATASET.imagesize,
        split     = DatasetSplit.TRAIN,
    )

    testset = create_dataset(
        dataname  = cfg.DATASET.name,
        source    = cfg.DATASET.datadir,
        classname = cfg.DATASET.classname,
        resize    = cfg.DATASET.resize,
        imagesize = cfg.DATASET.imagesize,
        split     = DatasetSplit.TEST,
    )

    # build dataloader
    trainloader = DataLoader(trainset, shuffle=False, batch_size=cfg.TRAIN.batch_size, num_workers=cfg.TRAIN.num_workers)
    testloader = DataLoader(testset, shuffle=False, batch_size=cfg.TRAIN.test_batch_size, num_workers=cfg.TRAIN.num_workers)

    # create model
    params = {}
    if cfg.MODEL.name == "HierarchicalPatchCore":
        params.update({
            "semantic_layer_to_extract_from" : cfg.MODEL.get("semantic_layer"),
            "semantic_batch_size"            : cfg.MODEL.get("semantic_batch_size"),
            "known_class"                    : cfg.MODEL.get("known_class", False),
            "clustering_method"              : cfg.MODEL.get("clustering_method", "finch"),
            "dataset_name"                   : cfg.DATASET.name,
            "n_clusters"                     : cfg.MODEL.get("n_clusters", None)
        })
        
    model = create_model(
        modelname                = cfg.MODEL.name,
        device                   = device,
        backbone                 = backbones.load(cfg.MODEL.backbone),
        layers_to_extract_from   = cfg.MODEL.layers,
        input_shape              = trainset.imagesize,
        pretrain_embed_dimension = cfg.MODEL.pretrained_embed_dim,
        target_embed_dimension   = cfg.MODEL.target_embed_dim,
        patchsize                = cfg.MODEL.patchsize,
        anomaly_scorer_num_nn    = cfg.MODEL.anomaly_nn,
        featuresampler           = get_sampler(name=cfg.MODEL.sampler, percentage=cfg.MODEL.sampler_ratio, device=device, seed=cfg.DEFAULT.seed, **cfg.MODEL.get('params', {})),
        nn_method                = common.FaissNN(on_gpu=False, num_workers=cfg.MODEL.faiss.num_workers),
        savedir                  = savedir,
        **params
    )

    if cfg.DEFAULT.get('test') == False:
        OmegaConf.save(cfg, os.path.join(savedir, 'config.yaml'))
        # build memory bank
        model.fit(trainloader)
        model.save_to_path(save_path=savedir)
    
    model.load_from_path(
        load_path     = savedir, 
        backbone_name = cfg.MODEL.backbone, 
        nn_method     = common.FaissNN(on_gpu=cfg.MODEL.faiss.use_gpu, num_workers=cfg.MODEL.faiss.num_workers)
    )
    
    # evaluation
    results, binary_results, eval_args, infer_time = evaluate(model, testloader)
    
    # save results
    json.dump(results, open(os.path.join(savedir, 'test_results.json'), 'w'), indent='\t')
    json.dump(binary_results, open(os.path.join(savedir, 'test_binary_results.json'), 'w'), indent='\t')
    pickle.dump(eval_args, open(os.path.join(savedir, 'eval_args.pkl'), 'wb'))
    json.dump({'infer_time': infer_time}, open(os.path.join(savedir, 'infer_time.json'), 'w'), indent='\t')
    
    
if __name__ == '__main__':
    args = OmegaConf.from_cli()
    # load default config
    cfg = OmegaConf.load(args.config)
    del args['config']
    
    # merge config with new keys
    cfg = OmegaConf.merge(cfg, args)
    
    print(OmegaConf.to_yaml(cfg))
    
    run(cfg)