Update pages/2_SmoothGrad.py

pages/2_SmoothGrad.py

@@ -16,40 +16,23 @@ BACKGROUND_COLOR = '#bcd0e7'
 
 
 st.title('Feature attribution with SmoothGrad')
-st.write('Which features are responsible for the current prediction? ')
+st.write("""> **Which features are responsible for the current prediction?**
+
+In machine learning, it is helpful to identify the significant features of the input (e.g., pixels for images) that affect the model's prediction. 
+If the model makes an incorrect prediction, we might want to determine which features contributed to the mistake. 
+To do this, we can generate a feature importance mask, which is a grayscale image with the same size as the original image. 
+The brightness of each pixel in the mask represents the importance of that feature to the model's prediction.
+
+There are various methods to calculate an image sensitivity mask for a specific prediction. 
+One simple way is to use the gradient of a class prediction neuron concerning the input pixels, indicating how the prediction is affected by small pixel changes. 
+However, this method usually produces a noisy mask. 
+To reduce the noise, the [SmoothGrad](https://arxiv.org/abs/1706.03825) technique is used, which adds Gaussian noise to multiple copies of the image and averages the resulting gradients.
+""")
 
 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 = []
@@ -90,18 +73,14 @@ for i, model_name in enumerate(selected_models):
 
 
 # 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)
+if run_button:
+    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)
 
 
 @st.cache(allow_output_mutation=True)