ariG23498 commited on
Commit
3010861
1 Parent(s): 9d17ea5

chore: adding spinners and caching

Browse files
Files changed (1) hide show
  1. app.py +7 -4
app.py CHANGED
@@ -22,15 +22,18 @@ image, preprocessed_image = utils.load_image_from_url(
22
  )
23
  st.image(image, caption="Original Image")
24
 
25
- st.title("Attention Heat Maps")
26
-
27
  # Load the DINO model
28
- with st.spinner('Processing the heat maps. Wait for it 🔥...'):
29
- dino = from_pretrained_keras("probing-vits/vit-dino-base16")
 
 
 
30
 
 
31
  # Get the attention scores
32
  _, attention_score_dict = dino.predict(preprocessed_image)
33
 
 
34
  # De-normalize the image for visual clarity.
35
  in1k_mean = tf.constant([0.485 * 255, 0.456 * 255, 0.406 * 255])
36
  in1k_std = tf.constant([0.229 * 255, 0.224 * 255, 0.225 * 255])
 
22
  )
23
  st.image(image, caption="Original Image")
24
 
 
 
25
  # Load the DINO model
26
+ @st.cache(persist=True)
27
+ def load_dino():
28
+ return from_pretrained_keras("probing-vits/vit-dino-base16")
29
+
30
+ dino = load_dino()
31
 
32
+ with st.spinner("Generating the attention scores..."):
33
  # Get the attention scores
34
  _, attention_score_dict = dino.predict(preprocessed_image)
35
 
36
+ with st.spinner("Generating the heat maps..."):
37
  # De-normalize the image for visual clarity.
38
  in1k_mean = tf.constant([0.485 * 255, 0.456 * 255, 0.406 * 255])
39
  in1k_std = tf.constant([0.229 * 255, 0.224 * 255, 0.225 * 255])