Update app

app.py

@@ -411,8 +411,8 @@ def molecules_from_image():
         print((top_probs, top_labels))
 
 def images_from_molecule():
-    #st.markdown("Enter a query molecule in [SMILES](https://en.wikipedia.org/wiki/Simplified_molecular-input_line-entry_system) format",)
-    smiles = st.text_input("Enter a query molecule in [SMILES](https://en.wikipedia.org/wiki/Simplified_molecular-input_line-entry_system) format", value="CC(=O)OC1=CC=CC=C1C(=O)O", placeholder="CC(=O)OC1=CC=CC=C1C(=O)O")
+    st.markdown("Enter a query molecule in [SMILES](https://en.wikipedia.org/wiki/Simplified_molecular-input_line-entry_system) format",)
+    smiles = st.text_input("Enter a query molecule in SMILES format", value="CC(=O)OC1=CC=CC=C1C(=O)O", placeholder="CC(=O)OC1=CC=CC=C1C(=O)O", label_visibility="collapsed")
     if smiles:
         smiles = [smiles]
         morgan = [morgan_from_smiles(s) for s in smiles]
@@ -436,6 +436,10 @@ def images_from_molecule():
         with col3:
             st.write("")
 
+        top_n = st.selectbox(
+            "How many images would you like to be retrieve?",
+            ("5", "10", "20")
+        )
 
         img_features_torch = torch.load(image_features, map_location=device)
         img_features = img_features_torch["img_features"]
@@ -443,7 +447,7 @@ def images_from_molecule():
 
         logits = mol_features @ img_features.T
         img_probs = (30.0 * logits).softmax(dim=-1)
-        top_probs, top_labels = img_probs.cpu().topk(5, dim=-1)
+        top_probs, top_labels = img_probs.cpu().topk(top_n, dim=-1)
 
         top_probs = torch.flatten(top_probs)
         top_labels = torch.flatten(top_labels)
@@ -454,6 +458,7 @@ def images_from_molecule():
 
         images_dict = np.load(images_arr, allow_pickle = True)
 
+
         st.write("Retrieved images from the Cell Painting database")
         with st.container():
             columns = st.columns(len(top_probs))