fix class order in explorer / vis deprec. warnings

Author: mluerig

bioencoder/scripts/interactive_plots.py

@@ -77,7 +77,7 @@ def interactive_plots(
     plot_dir = os.path.join(root_dir, "plots", run_name)
     os.makedirs(plot_dir, exist_ok=True)
     plot_path = os.path.join(plot_dir, "embeddings_interactive_plot.html")
-    if not overwrite and not (return_embeddings or return_coords):
+    if not overwrite and not return_results:
         assert not os.path.isfile(plot_path), f"File already exists: {plot_path}"
 
     ## Load model and set up

bioencoder/scripts/model_explorer.py

@@ -9,6 +9,7 @@
 import torch
 import streamlit as st
 from streamlit_option_menu import option_menu
+from torchvision.datasets import ImageFolder
 
 from bioencoder import config, utils, vis
 
@@ -63,9 +64,7 @@ def model_explorer(
     ## load bioencoer config
     root_dir = config.root_dir
     run_name = config.run_name
-    
-    class_names = os.listdir(os.path.join(root_dir, "data", run_name, "train"))
-
+        
     ## load config
     hyperparams = utils.load_yaml(config_path)
 
@@ -79,7 +78,6 @@ def model_explorer(
 
     ## get swa path
     ckpt_pretrained = os.path.join(root_dir, "weights", run_name, stage, "swa")
-    
     if stage == 'first': 
         vis_funcs = ['Filters', 'Activations', 'Saliency']
     else:
@@ -90,28 +88,30 @@ def model_explorer(
 
     # Sidebar    
     img_path = "https://github.com/agporto/BioEncoder/raw/main/assets/bioencoder_logo.png"
-    st.sidebar.image(img_path, use_column_width=True)
+    st.sidebar.image(img_path, width='stretch')
     st.sidebar.title("BioEncoder Model Explorer")
 
     # Image upload
     uploaded_file = st.sidebar.file_uploader("Upload an Image", type=["png", "jpg", "jpeg"])
 
-    ## get image transformations
-    transform = utils.get_transforms(hyperparams, no_aug=True)
-
     # Load the model and add to cache
     model = load_model(
         ckpt_pretrained, 
         backbone, 
         num_classes, 
         stage
         )
-
+    
+    ## get class names
+    train_folder = os.path.join(root_dir, "data", run_name, "train")
+    train_folder_timm = ImageFolder(train_folder)
+    class_names = train_folder_timm.classes         
+                             
     if uploaded_file is not None:
 
         # Display the uploaded image
         image = Image.open(uploaded_file).convert('RGB')
-        st.sidebar.image(image, caption="Input Image", use_column_width=True)
+        st.sidebar.image(image, caption="Input Image", width='stretch')
 
         # resize image
         image_resized = image.resize((img_size, img_size))

bioencoder/scripts/train.py

@@ -84,9 +84,11 @@ def train(
         "valid_batch_size": hyperparams["dataloaders"]["valid_batch_size"],
     }
     num_workers = hyperparams["dataloaders"]["num_workers"]
-    aug_sample = hyperparams["augmentations"].get("sample_save", False)
-    aug_sample_n = hyperparams["augmentations"].get("sample_n", 5)
-    aug_sample_seed = hyperparams["augmentations"].get("sample_seed", 42)
+    aug_config = hyperparams.get("augmentations", {})
+    aug_sample = aug_config.get("sample_save", False)
+    aug_sample_n = aug_config.get("sample_n", 5)
+    aug_sample_seed = aug_config.get("sample_seed", 42)
+
 
     ## manage directories and paths
     data_dir = os.path.join(root_dir, "data", run_name)