Spaces:
Runtime error
Runtime error
File size: 5,470 Bytes
1cf3825 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 | import argparse
import os
import sys
from pathlib import Path
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
root = Path(__file__).resolve().parents[1]
sys.path.append(str(root))
from src.data import get_datasets, prepare_dataset
from src.models import get_model
from src.utils import make_gradcam_heatmap, overlay_heatmap
def parse_args():
parser = argparse.ArgumentParser(description='Generate explainability outputs for brain tumor models')
parser.add_argument('--model', choices=['cnn', 'transfer', 'vit'], default='cnn')
parser.add_argument('--dataset', default='dataset')
parser.add_argument('--weights', required=True)
parser.add_argument('--batch_size', type=int, default=8)
parser.add_argument('--output', default='artifacts')
parser.add_argument('--examples', type=int, default=4)
return parser.parse_args()
def _get_default_conv_layer(model_type):
if model_type == 'cnn':
return 'conv_block_3'
if model_type == 'transfer':
return 'conv5_block3_out'
return None
def _get_sample_images(dataset, max_examples=4):
images = []
labels = []
for batch, (x, y) in enumerate(dataset):
for i in range(x.shape[0]):
if len(images) >= max_examples:
return np.array(images), np.array(labels)
images.append(x[i].numpy())
labels.append(int(y[i].numpy()))
if len(images) >= max_examples:
break
return np.array(images), np.array(labels)
def _plot_image(image, title, save_path):
plt.figure(figsize=(5, 5))
plt.imshow(image.astype('uint8'))
plt.title(title)
plt.axis('off')
plt.tight_layout()
plt.savefig(save_path)
plt.close()
def vit_patch_saliency(model, image, image_size=(224, 224)):
"""Saliency for the hybrid ResNet50+ViT classifier defined in src/models.py.
The hybrid model projects the ResNet50 feature map (7x7) into patch tokens via
a 1x1 Conv2D named 'hybrid_patch_projection', then reshapes to a sequence and
passes it through transformer blocks. We compute gradient-based saliency on
the patch token sequence (after position embedding) and reshape it back to
a 7x7 grid before resizing to the input resolution.
"""
try:
token_layer = model.get_layer('hybrid_patch_tokens')
except ValueError:
token_layer = model.get_layer('hybrid_patch_projection')
token_model = tf.keras.Model(inputs=model.inputs, outputs=token_layer.output)
image_batch = tf.expand_dims(image, axis=0)
with tf.GradientTape() as tape:
tokens = token_model(image_batch)
tape.watch(tokens)
predictions = model(image_batch)
loss = predictions[:, 0]
grads = tape.gradient(loss, tokens)
importance = tf.reduce_mean(tf.abs(grads * tokens), axis=-1)
importance = tf.squeeze(importance).numpy()
num_tokens = importance.shape[0] if importance.ndim == 1 else importance.size
side = int(round(num_tokens ** 0.5))
if side * side != num_tokens:
side = max(1, int(np.floor(num_tokens ** 0.5)))
importance = importance[: side * side]
importance = importance.reshape(side, side)
importance = importance / (importance.max() + 1e-8)
importance = tf.expand_dims(importance, axis=-1)
importance = tf.image.resize(importance, image_size, method='bilinear').numpy()
return np.squeeze(importance)
def explain_examples(model, model_type, images, labels, output_dir):
os.makedirs(output_dir, exist_ok=True)
conv_layer = _get_default_conv_layer(model_type)
for idx, (image, label) in enumerate(zip(images, labels)):
title = f'Label={label}'
sample_image = image.astype('uint8') if image.dtype != 'uint8' else image
_plot_image(sample_image, f'Input {idx} ({title})', os.path.join(output_dir, f'input_{idx}.png'))
if model_type in ['cnn', 'transfer']:
heatmap = make_gradcam_heatmap(tf.expand_dims(image, axis=0), model, conv_layer)
overlay = overlay_heatmap(sample_image, heatmap)
_plot_image(overlay, f'Grad-CAM {idx} ({title})', os.path.join(output_dir, f'gradcam_{idx}.png'))
if model_type == 'vit':
heatmap = vit_patch_saliency(model, image)
overlay = overlay_heatmap(sample_image, heatmap)
_plot_image(overlay, f'ViT Patch Saliency {idx} ({title})', os.path.join(output_dir, f'vit_saliency_{idx}.png'))
prediction = model.predict(tf.expand_dims(image, axis=0), verbose=0)[0][0]
print(f'Example {idx}: true={label}, score={prediction:.4f}')
def main():
args = parse_args()
train_ds, val_ds, test_ds = get_datasets(args.dataset, batch_size=args.batch_size)
eval_ds = test_ds if test_ds is not None else val_ds
if eval_ds is None:
raise ValueError('No validation or test dataset available for explanation.')
eval_ds = prepare_dataset(eval_ds)
model = get_model(args.model)
model.load_weights(args.weights)
images, labels = _get_sample_images(eval_ds, max_examples=args.examples)
explain_dir = os.path.join(args.output, args.model, 'explain')
explain_examples(model, args.model, images, labels, explain_dir)
print(f'Explainability outputs saved to {explain_dir}')
if __name__ == '__main__':
main()
|