wasserstein_distance.py

import os
import json
import pandas as pd
import numpy as np
from tqdm.auto import tqdm
from omegaconf import OmegaConf

from torch.utils.data import  DataLoader

import pandas as pd

from scipy.stats import wasserstein_distance

import common
from models import create_model
import backbones
import utils
from dataset import DatasetSplit, create_dataset
from main import get_sampler
from metrics import score_scaling


class_info = {
    'MVTec':{
        'objects' : [
            "bottle",
            "cable",
            "capsule",
            "hazelnut",
            "metal_nut",
            "pill",
            "screw",
            "toothbrush",
            "transistor",
        ],
        'textures' : [
            "zipper",
            "carpet",
            "grid",
            "leather",
            "tile",
            "wood",
        ]
    },
    'VisA':{
        'complex' : [
            'pcb1',
            'pcb2',
            'pcb3',
            'pcb4',
        ],
        'multi' : [
            'macaroni1',
            'macaroni2',
            'capsules',
            'candle',
        ],
        'single' : [
            'cashew',
            'chewinggum',
            'fryum',
            'pipe_fryum'
        ]
    },
    'MPDD': {
        'objects' : [
            'bracket_black', 
            'bracket_brown', 
            'bracket_white', 
            'connector', 
            'metal_plate', 
            'tubes'
        ]
    }
}
    

def get_results_df(savedir):
    cfg = OmegaConf.load(os.path.join(savedir, 'config.yaml'))
    
    # device
    device = utils.set_torch_device(gpu_ids=cfg.DEFAULT.device_ids)

    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
    testloader = DataLoader(testset, shuffle=False, batch_size=cfg.TRAIN.test_batch_size, num_workers=cfg.TRAIN.num_workers)

    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"),
        })
        
    model = create_model(
        modelname                = cfg.MODEL.name,
        device                   = device,
        backbone                 = backbones.load(cfg.MODEL.backbone),
        layers_to_extract_from   = cfg.MODEL.layers,
        input_shape              = testset.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=cfg.MODEL.faiss.use_gpu, num_workers=cfg.MODEL.faiss.num_workers),
        **params
    )

    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)
    )

    # prediction
    outputs = model.predict(
        testloader
    )

    scores = outputs['scores']

    # score and segmentation scaling
    if 'semantic_match_indices' in outputs:
        semantic_match_indices = outputs['semantic_match_indices']
        
        scores_scaled = np.empty_like(scores)

        for idx in np.unique(semantic_match_indices):
            match_idx = np.where(semantic_match_indices==idx)[0]
            scores_scaled[match_idx] = score_scaling(scores=np.array(scores)[match_idx])
    
    else:
        scores_scaled = score_scaling(scores=np.array(scores))
    
    anomaly_labels = [
        x[1] != "good" for x in testloader.dataset.data_to_iterate
    ]

    df = pd.DataFrame({
        'scores'    : scores_scaled, 
        'labels'    : anomaly_labels, 
        'class'     : [x[0] for x in testloader.dataset.data_to_iterate],
        'mask_path' : [x[-1] for x in testloader.dataset.data_to_iterate],
    })

    return df

def save_wd_score(df, savedir):
    wd_scores = {}
    for cls in df['class'].unique():
        df_cls = df[df['class'] == cls]
        
        sample1 = df_cls[df_cls['labels']==True]['scores'].values
        sample2 = df_cls[df_cls['labels']==False]['scores'].values

        wass_distance = wasserstein_distance(sample1, sample2)

        wd_scores[cls] = wass_distance
        
    # save wasserstein distance
    with open(os.path.join(savedir, 'wd_scores.json'), 'w') as f:
        json.dump(wd_scores, f, indent='\t')



if __name__ == "__main__":
    args = OmegaConf.from_cli()
    
    for p_cls in tqdm(os.listdir(args.savedir)):
        savedir_cls = os.path.join(args.savedir, p_cls)
            
        df = get_results_df(savedir=savedir_cls)

        save_wd_score(df=df, savedir=savedir_cls)
        df.to_csv(os.path.join(savedir_cls, 'anomaly_score.csv'), index=False)