Update app.py

app.py

@@ -3,6 +3,7 @@ import torch
 from transformers import ViTFeatureExtractor, ViTModel
 from skops import hub_utils
 from einops import reduce
+from torchvision.transforms.functional import to_pil_image
 import matplotlib.pyplot as plt
 import seaborn as sns
 
@@ -57,14 +58,18 @@ def classify_and_heatmap(input_img):
   ax_orig_img = fig.add_subplot(gs[:, 0])
 
   # plot original image
-  img = feature_extractor.to_pil_image(
+  img = to_pil_image(
       inputs['pixel_values'].squeeze(0) * torch.tensor(feature_extractor.image_std).view(-1, 1, 1) + torch.tensor(feature_extractor.image_mean).view(-1, 1, 1)
   )
   ax_orig_img.imshow(img)
   ax_orig_img.axis('off')
 
   # plot patch contributions
-  patch_contributions = logistic_regression.coef_ @ patch_embeddings.T.numpy()
+  patch_contributions = (
+        logistic_regression.coef_ \
+        @ patch_embeddings.T.numpy() \
+        + logistic_regression.intercept_.reshape(-1, 1) / (num_patches_side ** 2)
+    ).reshape(-1, num_patches_side, num_patches_side)
   vmin = patch_contributions.min()
   vmax = patch_contributions.max()