Spaces:
Build error
Build error
File size: 5,317 Bytes
18f2f54 0c1e42b 18f2f54 0c1e42b 18f2f54 14d54e2 0c1e42b 18f2f54 0c1e42b 18f2f54 |
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 |
import streamlit as st
import pandas as pd
import numpy as np
import random
from backend.utils import make_grid, load_dataset, load_model, load_images
from backend.smooth_grad import generate_smoothgrad_mask, ShowImage, fig2img
from transformers import AutoFeatureExtractor, AutoModelForImageClassification
import torch
from matplotlib.backends.backend_agg import RendererAgg
_lock = RendererAgg.lock
st.set_page_config(layout='wide')
BACKGROUND_COLOR = '#bcd0e7'
st.title('Feature attribution with SmoothGrad')
st.write('Which features are responsible for the current prediction? ')
imagenet_df = pd.read_csv('./data/ImageNet_metadata.csv')
# --------------------------- LOAD function -----------------------------
# @st.cache(allow_output_mutation=True)
# @st.cache_data
# def load_images(image_ids):
# images = []
# for image_id in image_ids:
# dataset = load_dataset(image_id//10000)
# images.append(dataset[image_id%10000])
# return images
# @st.cache(allow_output_mutation=True, suppress_st_warning=True, show_spinner=False)
# @st.cache_resource
# def load_model(model_name):
# with st.spinner(f"Loading {model_name} model! This process might take 1-2 minutes..."):
# if model_name == 'ResNet':
# model_file_path = 'microsoft/resnet-50'
# feature_extractor = AutoFeatureExtractor.from_pretrained(model_file_path, crop_pct=1.0)
# model = AutoModelForImageClassification.from_pretrained(model_file_path)
# model.eval()
# elif model_name == 'ConvNeXt':
# model_file_path = 'facebook/convnext-tiny-224'
# feature_extractor = AutoFeatureExtractor.from_pretrained(model_file_path, crop_pct=1.0)
# model = AutoModelForImageClassification.from_pretrained(model_file_path)
# model.eval()
# else:
# model = torch.hub.load('pytorch/vision:v0.10.0', 'mobilenet_v2', pretrained=True)
# model.eval()
# feature_extractor = None
# return model, feature_extractor
images = []
image_ids = []
# INPUT ------------------------------
st.header('Input')
with st.form('smooth_grad_form'):
st.markdown('**Model and Input Setting**')
selected_models = st.multiselect('Model', options=['ConvNeXt', 'ResNet', 'MobileNet'])
selected_image_set = st.selectbox('Image set', ['User-defined set', 'Random set'])
summit_button = st.form_submit_button('Set')
if summit_button:
setting_container = st.container()
# for id in image_ids:
# images = load_images(image_ids)
with st.form('2nd_form'):
st.markdown('**Image set setting**')
if selected_image_set == 'Random set':
no_images = st.slider('Number of images', 1, 50, value=10)
image_ids = random.sample(list(range(50_000)), k=no_images)
else:
text = st.text_area('Specific Image IDs', value='0')
image_ids = list(map(lambda x: int(x.strip()), text.split(',')))
run_button = st.form_submit_button('Display output')
if run_button:
for id in image_ids:
images = load_images(image_ids)
st.header('Output')
models = {}
feature_extractors = {}
for i, model_name in enumerate(selected_models):
models[model_name], feature_extractors[model_name] = load_model(model_name)
# DISPLAY ----------------------------------
header_cols = st.columns([1, 1] + [2]*len(selected_models))
header_cols[0].markdown(f'<div style="text-align: center;margin-bottom: 10px;background-color:{BACKGROUND_COLOR};"><b>Image ID</b></div>', unsafe_allow_html=True)
header_cols[1].markdown(f'<div style="text-align: center;margin-bottom: 10px;background-color:{BACKGROUND_COLOR};"><b>Original Image</b></div>', unsafe_allow_html=True)
for i, model_name in enumerate(selected_models):
header_cols[i + 2].markdown(f'<div style="text-align: center;margin-bottom: 10px;background-color:{BACKGROUND_COLOR};"><b>{model_name}</b></div>', unsafe_allow_html=True)
grids = make_grid(cols=2+len(selected_models)*2, rows=len(image_ids)+1)
# grids[0][0].write('Image ID')
# grids[0][1].write('Original image')
# for i, model_name in enumerate(selected_models):
# models[model_name], feature_extractors[model_name] = load_model(model_name)
@st.cache(allow_output_mutation=True)
# @st.cache_data
def generate_images(image_id, model_name):
j = image_ids.index(image_id)
image = images[j]['image']
return generate_smoothgrad_mask(
image, model_name,
models[model_name], feature_extractors[model_name], num_samples=10)
with _lock:
for j, (image_id, image_dict) in enumerate(zip(image_ids, images)):
grids[j][0].write(f'{image_id}. {image_dict["label"]}')
image = image_dict['image']
ori_image = ShowImage(np.asarray(image))
grids[j][1].image(ori_image)
for i, model_name in enumerate(selected_models):
# ori_image, heatmap_image, masked_image = generate_smoothgrad_mask(image,
# model_name, models[model_name], feature_extractors[model_name], num_samples=10)
heatmap_image, masked_image = generate_images(image_id, model_name)
# grids[j][1].image(ori_image)
grids[j][i*2+2].image(heatmap_image)
grids[j][i*2+3].image(masked_image) |