Update app.py

app.py

@@ -27,10 +27,10 @@ with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
 
 st.title('Find my pic!')
 
-def find_image_disc(prompt, df):
+def find_image_disc(prompt, df, top_k):
     img_descs = []
     img_descs_vit = []
-    list_images_names, list_images_names_vit = get_similiarity(prompt, model_resnet, model_vit, 3)
+    list_images_names, list_images_names_vit = get_similiarity(prompt, model_resnet, model_vit, top_k)
     for img in list_images_names:
         img_descs.append(random.choice(df[df['image_name'] == img.split('/')[-1]]['comment'].values).replace('.', ''))
     #vit
@@ -40,7 +40,7 @@ def find_image_disc(prompt, df):
     return list_images_names, img_descs, list_images_names_vit, img_descs_vit
 
 txt = st.text_area("Describe the picture you'd like to see")
-
+top_k = st.slider('Number of images', 1, 5, 3)
 df = pd.read_csv('results.csv', 
                  sep = '|', 
                  names = ['image_name', 'comment_number', 'comment'], 
@@ -49,8 +49,7 @@ df = pd.read_csv('results.csv',
 
 if txt is not None:
      if st.button('Find!'):
-
-        list_images, img_desc, list_images_vit, img_descs_vit = find_image_disc(txt, df)
+        list_images, img_desc, list_images_vit, img_descs_vit = find_image_disc(txt, df, top_k)
         col1, col2 = st.columns(2)
         col1.header('ResNet50')
         col2.header('ViT 32')