Allow model selection

app.py

@@ -43,7 +43,6 @@ imgname = "I1"
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model_type = "RN50"
-image_resolution = 520
 
 ######### CLOOME FUNCTIONS #########
 def convert_models_to_fp32(model):
@@ -410,9 +409,11 @@ def molecules_from_image(top_n):
         print(mol_probs.sum(dim=-1))
         print((top_probs, top_labels))
 
-def images_from_molecule(top_n):
+def images_from_molecule(top_n, model_path):
     #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")
+    
+    
     if smiles:
         smiles = [smiles]
         morgan = [morgan_from_smiles(s) for s in smiles]
@@ -423,7 +424,7 @@ def images_from_molecule(top_n):
         fps_fname = save_hdf(morgan, molnames, molpath)
         mol_imgs = draw_molecules(smiles)
 
-        mol_features, mol_ids = main(mol_index, MODEL_PATH, model_type, mol_path=molpath, image_resolution=image_resolution)
+        mol_features, mol_ids = main(mol_index, model_path, model_type, mol_path=molpath, image_resolution=image_resolution)
 
         col1, col2, col3 = st.columns(3)
 
@@ -494,5 +495,27 @@ n_objects = st.sidebar.selectbox(
             "How many objects would you like to retrieve?",
             ("5", "10", "20"))
 
-page_names_to_funcs[selected_page](n_objects)
+
+selected_model = st.sidebar.selectbox(
+                "Select a CLOOME model to load",
+                ("CLOOME (default)", "CLOOME (batch size 128)", "CLOOME (fullres)", "CLOOME (imgres 320)"))
+
+model_dict = {
+    "CLOOME (default)" : "cloome_default.pt",
+    "CLOOME (CLIP imgres 320)" : "cloome_cli_imres320.pt",
+    "CLOOME (fullres)" : "cloome_fullres.pt",
+    "CLOOME (CLOOB imgres 320)" : "cloome_imres320.pt"
+
+}
+
+model_file = model_dict[selected_model]
+model_path = os.path.join(datapath, model_file)
+
+if model_path.endswith("320.pt"):
+    image_resolution = 320
+else: 
+    image_resolution = 520
+
+
+page_names_to_funcs[selected_page](n_objects, selected_model)
 

data/cloome_cli_imres320.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:91d042823294919768aae7bf2a489c6c6e0c9535b8adc1be19adea8fce03cb5a
+size 352014131

data/cloome_default.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:39e0c98d47b18ce913f4bcb1a1bc89d26ca9938ee74646c32a46ad236cddbc38
+size 352014131

data/cloome_fullres.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f412fe015caef7d9b9012a7ab8885081b089109abbbee76f5c5c609935151ec5
+size 352013623

data/cloome_imres320.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5eee478ca947ae9708d5432d427fb086e2141f1b0a81d6d9e8fe11fc4f29a1da
+size 352013623