import os from concurrent.futures import ThreadPoolExecutor, as_completed import clip import h5py import ml_collections import numpy as np import open_clip import streamlit as st import torch from huggingface_hub import hf_hub_download from app_lib.ckde import cKDE from app_lib.utils import SUPPORTED_MODELS from ibydmt.test import xSKIT rng = np.random.default_rng() def _get_open_clip_model(model_name, device): backbone = model_name.split(":")[-1] model, _, preprocess = open_clip.create_model_and_transforms( SUPPORTED_MODELS[model_name], device=device ) model.eval() tokenizer = open_clip.get_tokenizer(backbone) return model, preprocess, tokenizer def _get_clip_model(model_name, device): backbone = model_name.split(":")[-1] model, preprocess = clip.load(backbone, device=device) tokenizer = clip.tokenize return model, preprocess, tokenizer def _load_model(model_name, device): if "open_clip" in model_name: model, preprocess, tokenizer = _get_open_clip_model(model_name, device) elif "clip" in model_name: model, preprocess, tokenizer = _get_clip_model(model_name, device) return model, preprocess, tokenizer @torch.no_grad() @torch.cuda.amp.autocast() def _encode_concepts(tokenizer, model, concepts, device): concepts_text = tokenizer(concepts).to(device) concept_features = model.encode_text(concepts_text) concept_features /= torch.linalg.norm(concept_features, dim=-1, keepdim=True) return concept_features.cpu().numpy() @torch.no_grad() @torch.cuda.amp.autocast() def _encode_image(model, preprocess, image, device): image = preprocess(image) image = image.unsqueeze(0) image = image.to(device) image_features = model.encode_image(image) image_features /= image_features.norm(dim=-1, keepdim=True) return image_features.cpu().numpy() @torch.no_grad() @torch.cuda.amp.autocast() def _encode_class_name(tokenizer, model, class_name, device): class_text = tokenizer([f"A photo of a {class_name}"]).to(device) class_features = model.encode_text(class_text) class_features /= torch.linalg.norm(class_features, dim=-1, keepdim=True) return class_features.cpu().numpy() def _load_dataset(dataset_name, model_name): dataset_path = hf_hub_download( repo_id="jacopoteneggi/IBYDMT", filename=f"{dataset_name}_{model_name}_train.h5", repo_type="dataset", ) with h5py.File(dataset_path, "r") as dataset: embedding = dataset["embedding"][:] return embedding def _sample_random_subset(concept_idx, concepts, cardinality): sample_idx = list(set(range(len(concepts))) - {concept_idx}) return rng.permutation(sample_idx)[:cardinality].tolist() def _test(testing_config, z, concept_idx, concepts, cardinality, sampler, classifier): def cond_p(z, cond_idx, m): _, sample_h = sampler.sample(z, cond_idx, m=m) return sample_h def f(h): output = h @ classifier.T return output.squeeze() rejected_hist, tau_hist, wealth_hist, subset_hist = [], [], [], [] for _ in range(testing_config.r): subset_idx = _sample_random_subset(concept_idx, concepts, cardinality) subset = [concepts[idx] for idx in subset_idx] tester = xSKIT(testing_config) rejected, tau = tester.test( z, concept_idx, subset_idx, cond_p, f, interrupt_on="max_wealth", max_wealth=3 * 1 / testing_config.significance_level, ) wealth = tester.wealth._wealth wealth = wealth + [wealth[-1]] * (testing_config.tau_max - len(wealth)) rejected_hist.append(rejected) tau_hist.append(tau) wealth_hist.append(wealth) subset_hist.append(subset) return { "concept": concepts[concept_idx], "rejected": rejected_hist, "tau": tau_hist, "wealth": wealth_hist, "subset": subset_hist, } def get_testing_config(**kwargs): testing_config = st.session_state.testing_config = ml_collections.ConfigDict() testing_config.significance_level = kwargs.get("significance_level", 0.05) testing_config.wealth = kwargs.get("wealth", "ons") testing_config.bet = kwargs.get("bet", "tanh") testing_config.kernel = kwargs.get("kernel", "rbf") testing_config.kernel_scale_method = kwargs.get("kernel_scale_method", "quantile") testing_config.kernel_scale = kwargs.get("kernel_scale", 0.5) testing_config.tau_max = kwargs.get("tau_max", 200) testing_config.r = kwargs.get("r", 10) return testing_config def load_precomputed_results(image_name): results = np.load( os.path.join("assets", "results", f"{image_name.split('.')[0]}.npy"), allow_pickle=True, ).item() return results def test( testing_config, image, class_name, concepts, cardinality, dataset_name, model_name, device=torch.device("cuda" if torch.cuda.is_available() else "cpu"), with_streamlit=True, ): if with_streamlit: with st.spinner("Loading model"): model, preprocess, tokenizer = _load_model(model_name, device) else: model, preprocess, tokenizer = _load_model(model_name, device) if with_streamlit: with st.spinner("Encoding concepts"): cbm = _encode_concepts(tokenizer, model, concepts, device) else: cbm = _encode_concepts(tokenizer, model, concepts, device) if with_streamlit: with st.spinner("Encoding image"): h = _encode_image(model, preprocess, image, device) else: h = _encode_image(model, preprocess, image, device) z = h @ cbm.T z = z.squeeze() if with_streamlit: progress_bar = st.progress( 0, text=( "Testing concepts (can take up to a minute) [0 /" f" {len(concepts)} completed]" ), ) progress_bar.progress( 1 / (len(concepts) + 1), text=( "Testing concepts (can take up to a minute) [0 /" f" {len(concepts)} completed]" ), ) embedding = _load_dataset(dataset_name, model_name) semantics = embedding @ cbm.T sampler = cKDE(embedding, semantics) classifier = _encode_class_name(tokenizer, model, class_name, device) with ThreadPoolExecutor() as executor: futures = [ executor.submit( _test, testing_config, z, concept_idx, concepts, cardinality, sampler, classifier, ) for concept_idx in range(len(concepts)) ] results = [] for idx, future in enumerate(as_completed(futures)): results.append(future.result()) if with_streamlit: progress_bar.progress( (idx + 2) / (len(concepts) + 1), text=( f"Testing concepts (can take up to a minute) [{idx + 1} /" f" {len(concepts)} completed]" ), ) rejected = np.empty((testing_config.r, len(concepts))) tau = np.empty((testing_config.r, len(concepts))) wealth = np.empty((testing_config.r, testing_config.tau_max, len(concepts))) for _results in results: concept_idx = concepts.index(_results["concept"]) rejected[:, concept_idx] = np.array(_results["rejected"]) tau[:, concept_idx] = np.array(_results["tau"]) wealth[:, :, concept_idx] = np.array(_results["wealth"]) tau /= testing_config.tau_max results = { "significance_level": testing_config.significance_level, "concepts": concepts, "rejected": rejected, "tau": tau, "wealth": wealth, } return results