"""Predicts the output of the model on the test data."""
import argparse
import os.path as op
import nibabel as nib
import numpy as np
import pandas as pd
import torch
from nilearn.image import load_img, resample_to_img
from braindec.cogatlas import CognitiveAtlas
from braindec.embedding import ImageEmbedding
from braindec.model import build_model
from braindec.utils import _get_device, _read_vocabulary, get_data_dir, images_have_same_fov
[docs]
def preprocess_image(
image,
standardize=False,
data_dir=None,
space="MNI152",
density=None,
image_emb_gene=None,
):
"""
Preprocess the image.
Args:
image: Images
"""
data_dir = get_data_dir(data_dir)
nilearn_dir = op.join(data_dir, "nilearn")
if image_emb_gene is None:
image_emb_gene = ImageEmbedding(
standardize=standardize,
nilearn_dir=nilearn_dir,
space=space,
density=density,
)
image_embedding_arr = image_emb_gene(image)
return torch.from_numpy(image_embedding_arr).float()
[docs]
def image_to_labels(
image,
model_path,
vocabulary,
vocabulary_emb,
prior_probability,
topk=10,
logit_scale=None,
return_posterior_probability=False,
device=None,
model=None,
**kwargs,
):
"""Predict the labels of an image using a pre-trained model."""
if device is None:
device = _get_device()
if isinstance(image, str):
image = load_img(image)
image_input = preprocess_image(image, **kwargs).to(device)
text_inputs = torch.from_numpy(vocabulary_emb).float().to(device)
prior_probability = torch.from_numpy(prior_probability).float().to(device)
# Normalize the embeddings
text_inputs = text_inputs / (text_inputs.norm(dim=-1, keepdim=True) + 1e-8)
image_input = image_input / (image_input.norm(dim=-1, keepdim=True) + 1e-8)
# Calculate features
model = build_model(model_path, device=device) if model is None else model
with torch.no_grad():
image_features, text_features = model(image_input, text_inputs) # normalized
# Get the scaling factor: inverse temperature
logit_scale = model.logit_scale.item() if logit_scale is None else logit_scale
# Pick the top topk most similar labels for the image
# similarity = logit_scale * image_features @ text_features.T
similarity = logit_scale * image_features @ text_features.T
likelihood = similarity.softmax(dim=-1)
# Flatten the probability distribution, since image_features is a single image
similarity = similarity.flatten()
likelihood = likelihood.flatten() # P(A|T)
joint_probability = likelihood * prior_probability # P(A|T) * P(A)
total_probability = joint_probability.sum() # P(T)
posterior_probability = joint_probability / total_probability # P(T|A) = P(A|T) * P(A) / P(T)
# Calculate the strength of the evidence using the Bayes factor
posterrior_odds = posterior_probability / (1 - posterior_probability)
prior_odds = prior_probability / (1 - prior_probability)
bayes_factor = posterrior_odds / prior_odds
# Get top tasks
top_task_prob, top_indices = posterior_probability.topk(topk)
top_indices = top_indices.cpu().detach().numpy()
top_task_prob = top_task_prob.cpu().detach().numpy()
similarity = similarity.cpu().detach().numpy()
likelihood = likelihood.cpu().detach().numpy()
joint_probability = joint_probability.cpu().detach().numpy()
prior_probability = prior_probability.cpu().detach().numpy()
bayes_factor = bayes_factor.cpu().detach().numpy()
task_prob_df = pd.DataFrame(
{
"pred": np.array(vocabulary)[top_indices],
"prob": top_task_prob,
"similarity": similarity[top_indices],
"likelihood": likelihood[top_indices],
"prior_prob": prior_probability[top_indices],
"joint_prob": joint_probability[top_indices],
"bayes_factor": bayes_factor[top_indices],
}
)
if return_posterior_probability:
return task_prob_df, posterior_probability
return task_prob_df
[docs]
def image_to_labels_hierarchical(
image,
model_path,
vocabulary,
vocabulary_emb,
prior_probability,
cognitiveatlas,
topk=10,
logit_scale=None,
device=None,
model=None,
**kwargs,
):
"""Predict the label of an image."""
task_prob_df, posterior_probability = image_to_labels(
image,
model_path,
vocabulary,
vocabulary_emb,
prior_probability,
topk=topk,
logit_scale=logit_scale,
return_posterior_probability=True,
device=device,
model=model,
**kwargs,
)
# Get mapping and names from cognitive atlas object
concept_to_task_idxs = cognitiveatlas.concept_to_task_idxs
process_to_concept_idxs = cognitiveatlas.process_to_concept_idxs
concept_names = cognitiveatlas.concept_names
process_names = cognitiveatlas.process_names
# Calculate P(C|A) = 1 - Prod(1 - P(T|A))
concept_posterior_probability = torch.zeros(len(concept_names)) # Pre-allocate tensor
for c_i in range(len(concept_names)):
task_indices = concept_to_task_idxs[c_i]
concept_probability = 1 - (1 - posterior_probability[task_indices]).prod()
concept_posterior_probability[c_i] = concept_probability.cpu().detach()
process_posterior_probability = torch.zeros(len(process_names)) # Pre-allocate tensor
for p_i in range(len(process_names)):
concept_indices = process_to_concept_idxs[p_i]
process_probability = 1 - (1 - concept_posterior_probability[concept_indices]).prod()
process_posterior_probability[p_i] = process_probability.cpu().detach()
# Get the top k predictions
top_concepts, top_concept_indices = concept_posterior_probability.topk(topk)
top_concept_indices = top_concept_indices.cpu().detach().numpy()
top_concepts = top_concepts.cpu().detach().numpy()
top_processes, top_process_indices = torch.sort(process_posterior_probability, descending=True)
top_process_indices = top_process_indices.cpu().detach().numpy()
top_processes = top_processes.cpu().detach().numpy()
process_prob_df = pd.DataFrame(
{
"pred": np.array(process_names)[top_process_indices],
"prob": top_processes,
}
)
concept_prob_df = pd.DataFrame(
{
"pred": np.array(concept_names)[top_concept_indices],
"prob": top_concepts,
}
)
return task_prob_df, concept_prob_df, process_prob_df
def _get_parser():
parser = argparse.ArgumentParser(description="Run gradient-decoding workflow")
parser.add_argument(
"--image",
dest="image",
required=True,
help="Path to image file (e.g., NIfTI format).",
)
parser.add_argument(
"--model",
dest="model",
required=True,
help="Path to the pre-trained CLIP model.",
)
parser.add_argument(
"--vocabulary",
dest="vocabulary",
required=True,
help="Path to the vocabulary file.",
)
parser.add_argument(
"--vocabulary_emb",
dest="vocabulary_emb",
required=True,
help="Path to the vocabulary embedding file.",
)
parser.add_argument(
"--vocabulary_prior",
dest="vocabulary_prior",
required=True,
help="Path to the vocabulary prior file.",
)
parser.add_argument(
"--cognitiveatlas",
dest="cognitiveatlas",
required=False,
help="Path to the cognitive atlas object file.",
type=str,
default=None,
)
parser.add_argument(
"--hierarchical",
dest="hierarchical",
required=False,
default=False,
)
parser.add_argument(
"--reduced",
dest="reduced",
required=False,
default=True,
help="Whether to use the reduced task set.",
)
parser.add_argument(
"--mask",
dest="mask",
required=True,
help="Path to the mask file.",
)
parser.add_argument(
"--topk",
dest="topk",
type=int,
default=10,
help="Number of top predictions to return (default: 10).",
)
parser.add_argument(
"--logit_scale",
dest="logit_scale",
type=float,
default=10.0,
help="Logit scale for temperature scaling (default: None).",
)
parser.add_argument(
"--device",
dest="device",
default=None,
help="Device to use for computation (default: None). Possible values: cpu, mps, cuda.",
)
parser.add_argument(
"--output",
dest="output_dir",
required=False,
help="Path to the output directory. Defaults to the same directory as the input image.",
)
return parser
def _main(argv=None):
option = _get_parser().parse_args(argv)
image_fn = option.image
model_fn = option.model
vocabulary_fn = option.vocabulary
vocabulary_emb_fn = option.vocabulary_emb
vocabulary_prior_fn = option.vocabulary_prior
cognitiveatlas_path = option.cognitiveatlas
hierarchical = option.hierarchical
reduced = option.reduced
topk = option.topk
logit_scale = option.logit_scale
device = option.device
mask_fn = option.mask
output_dir = option.output_dir
if hierarchical and cognitiveatlas_path is None:
raise ValueError(
"Cognitive atlas files are required for hierarchical decoding. "
"Please provide the --cognitiveatlas path argument."
)
output_dir = op.dirname(image_fn) if output_dir is None else op.abspath(output_dir)
image_name = op.basename(image_fn).split(".")[0]
img = nib.load(image_fn)
mask_img = nib.load(mask_fn)
if not images_have_same_fov(img, mask_img):
img = resample_to_img(img, mask_img)
vocabulary, vocabulary_emb, vocabulary_prior = _read_vocabulary(
vocabulary_fn,
vocabulary_emb_fn,
vocabulary_prior_fn,
)
if hierarchical:
reduced_tasks_fn = op.join(cognitiveatlas_path, "reduced_tasks.csv")
reduced_tasks_df = pd.read_csv(reduced_tasks_fn) if reduced else None
cognitiveatlas = CognitiveAtlas(
data_dir=cognitiveatlas_path,
task_snapshot=op.join(cognitiveatlas_path, "task_snapshot-02-19-25.json"),
concept_snapshot=op.join(
cognitiveatlas_path, "concept_extended_snapshot-02-19-25.json"
),
reduced_tasks=reduced_tasks_df,
)
task_prob_df, concept_prob_df, process_prob_df = image_to_labels_hierarchical(
img,
model_fn,
vocabulary,
vocabulary_emb,
vocabulary_prior,
cognitiveatlas,
topk=topk,
logit_scale=logit_scale,
device=device,
)
task_prob_df.to_csv(op.join(output_dir, f"{image_name}_task_predictions.csv"), index=False)
concept_prob_df.to_csv(
op.join(output_dir, f"{image_name}_concept_predictions.csv"), index=False
)
process_prob_df.to_csv(
op.join(output_dir, f"{image_name}_process_predictions.csv"), index=False
)
else:
prob_df = image_to_labels(
img,
model_fn,
vocabulary,
vocabulary_emb,
vocabulary_prior,
topk=topk,
logit_scale=logit_scale,
return_posterior_probability=False,
device=device,
)
prob_df.to_csv(op.join(output_dir, f"{image_name}_predictions.csv"), index=False)
if __name__ == "__main__":
_main()