# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
import numpy as np
import pandas as pd
import torch
import os
import multiprocessing as mp
import seaborn as sns
import glob
import pylab
import json
import sys
from tqdm.notebook import tqdm
import gc
import h5py
import torchvision
sys.path.insert(0, os.path.abspath('..'))
from src.wrappers import get_wrapper
from src.datasets import get_dataset

# Set constants
attributes = "5_o_Clock_Shadow Arched_Eyebrows Attractive Bags_Under_Eyes Bald Bangs Big_Lips Big_Nose Black_Hair Blond_Hair Blurry Brown_Hair Bushy_Eyebrows Chubby Double_Chin Eyeglasses Goatee Gray_Hair Heavy_Makeup High_Cheekbones Male Mouth_Slightly_Open Mustache Narrow_Eyes No_Beard Oval_Face Pale_Skin Pointy_Nose Receding_Hairline Rosy_Cheeks Sideburns Smiling Straight_Hair Wavy_Hair Wearing_Earrings Wearing_Hat Wearing_Lipstick Wearing_Necklace Wearing_Necktie Young".split(" ")
datadir = '../data'
results_path = '../logs/bte'

# -

# # Beyond Trivial Explanations
# 1. Parse results from the BTE experiment
# 2. Plot non-trivial explanations against VGG similarity

# +
def process_results(path):
    i, path = path
    if i % 50 == 0:
        print(i)

    with open(os.path.join(os.path.dirname(path), 'exp_dict.json'), 'r') as infile:
        exp_dict = json.load(infile)
    with open(os.path.join(datadir, os.path.join(exp_dict['generator_path']), 'exp_dict.json'), 'r') as infile:
        cls_dict = json.load(infile)
        cls_dict['labels_path'] = 'celeba_meta_explanations/list_eval_unbiased.csv'
    att_idx = attributes.index(exp_dict["attribute"])
    ret = {}
    keys = ['logits', 'labels', 'oracle_preds_attack', 'oracle_preds_reconstruction']
    with h5py.File(path, 'r') as infile:
        for k in keys: 
            ret[k] = infile[k][..., att_idx]
        for k in ['latent_similarity', 'idx']:
            ret[k] = infile[k][:]
    
    dataset = get_dataset(datadir, 'val', cls_dict)
    ret['labels'] = dataset.y[ret['idx'], att_idx, None]
    ret.update(exp_dict)
    gc.collect()
    return ret

NUM_CPU = 8
# Parse results in parallel
with mp.Pool(NUM_CPU) as pool:
    ret = pool.map(process_results, enumerate(glob.glob(os.path.join(results_path, '*/results.h5'))))

# Compute final scores and save into dataframe
df = pd.DataFrame(ret)
df['logits'] = df['logits'].apply(lambda x: x > 0)
df['oracle_preds_reconstruction'] = df['oracle_preds_reconstruction'].apply(lambda x: x > 0)
df['oracle_preds_attack'] = df['oracle_preds_attack'].apply(lambda x: x > 0)
df['non-trivial'] = [((attack != oracle) & (oracle_attack == oracle)).sum() / ((attack != oracle)).sum() for attack, oracle, oracle_attack, label in zip(df['logits'].values, df['oracle_preds_reconstruction'].values, df['oracle_preds_attack'].values, df['labels'].values)]
df['latent_similarity'] = df['latent_similarity'].apply(np.nanmean)
display(df[['method', 'lr', 'non-trivial', 'score', 'latent_similarity', 'lasso_weight', 'reconstruction_weight', 'diversity_weight']])
df_backup = df.copy()
# -

# ## BTE PLOT

pylab.figure(figsize=(20,20))
for method in list(sorted(set(df['method']))):
    sns.kdeplot(df[df['method'] == method]['latent_similarity'], y=df[df['method'] == method]['non-trivial'], levels=[0.22])
    sns.scatterplot(x=[df[df['method'] == method]['latent_similarity'].values.mean()], y=[df[df['method'] == method]['non-trivial'].values.mean()], label=method)
pylab.title('BTE')
pylab.legend()

StarcoderdataJupyterScriptsDedupFiltered