Merge remote-tracking branch 'refs/remotes/origin/dev' into dev

Author: mluerig

bioencoder/core/augmentations.py

@@ -2,7 +2,7 @@
 import albumentations as A
 from albumentations import pytorch as AT
 
-def get_transforms(config, valid=False):
+def get_transforms(config, no_aug=False):
     """
     Return a transformation pipeline based on the provided configuration.
 
@@ -13,14 +13,16 @@ def get_transforms(config, valid=False):
     Returns:
         albumentations.core.composition.Compose: The image transformation pipeline.
     """
-    default_size = 224
-    img_size = config.get('img_size', default_size)
+    
+    img_size = config.get('img_size')
+    if img_size is None:
+        raise ValueError("config must include 'img_size'")
     config_aug = config.get('augmentations', {})
     aug = get_aug_from_config(config_aug.get('transforms', []))
 
     return A.Compose([
         A.Resize(img_size, img_size, always_apply=True),
-        A.NoOp() if valid else aug,
+        A.NoOp() if no_aug else aug,
         A.Normalize(),
         AT.ToTensorV2()
     ])

bioencoder/core/utils.py

@@ -188,15 +188,17 @@ def build_transforms(config):
 
     """
     train_transforms = get_transforms(config)
-    valid_transforms = get_transforms(config, valid=True)
+    valid_transforms = get_transforms(config, no_aug=True)
 
     return {
         "train_transforms": train_transforms,
         "valid_transforms": valid_transforms
     }
 
 
-def build_loaders(data_dir, transforms, batch_sizes, num_workers, second_stage=False, is_supcon=False):
+def build_loaders(data_dir, transforms, batch_sizes, num_workers, 
+                  second_stage=False, is_supcon=False,
+                  shuffle_train=True, drop_last=True):
     """
     Build data loaders for training and validation.
     
@@ -230,19 +232,19 @@ def build_loaders(data_dir, transforms, batch_sizes, num_workers, second_stage=F
     train_loader = torch.utils.data.DataLoader(
         train_features_dataset, 
         batch_size=batch_sizes['train_batch_size'], 
-        shuffle=True,
+        shuffle=shuffle_train,
         num_workers=num_workers, 
         pin_memory=True, 
-        drop_last=(batch_sizes['train_batch_size'] is not None)
+        drop_last=drop_last and batch_sizes['train_batch_size'] is not None
     )
-        
+
     valid_loader = torch.utils.data.DataLoader(
         valid_dataset, 
         batch_size=batch_sizes['valid_batch_size'], 
         shuffle=False,
         num_workers=num_workers, 
         pin_memory=True, 
-        drop_last=(batch_sizes['valid_batch_size'] is not None)
+        drop_last=drop_last
     )
 
     loaders = {

bioencoder/scripts/inference.py

@@ -85,7 +85,7 @@ def inference(
     utils.set_seed()
 
     ## get transformations
-    transform = utils.get_transforms(hyperparams, valid=False)
+    transform = utils.get_transforms(hyperparams, no_aug=True)
 
     ## build model
     if config.model_path != ckpt_pretrained:

bioencoder/scripts/interactive_plots.py

@@ -61,7 +61,7 @@ def interactive_plots(
         "valid_batch_size": hyperparams.get("dataloaders", {}).get("valid_batch_size",1),
     }
     num_workers = hyperparams.get("dataloaders", {}).get("num_workers", 4)
-    perplexity = hyperparams.get("perplexity", 30)
+    perplexity = hyperparams.get("perplexity")
 
     plot_config = {
         "color_classes": hyperparams.get("color_classes", None),
@@ -70,43 +70,40 @@ def interactive_plots(
         "point_size": hyperparams.get("point_size", 10),
     }
 
-    
-    ## Set up directories
+    ## directories and file management
     data_dir = os.path.join(root_dir, "data", run_name)
-    plot_path = os.path.join(root_dir, "plots", run_name, f"embeddings_{run_name}.html")
+    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 kwargs.get("ret_embeddings"):
-        assert not os.path.isfile(plot_path), f"File exists: {plot_path}"
+        assert not os.path.isfile(plot_path), f"File already exists: {plot_path}"
 
     ## Load model and set up
     print(f"Checkpoint: using {checkpoint} of {stage} stage")
     ckpt_pretrained = os.path.join(root_dir, "weights", run_name, stage, checkpoint)
     utils.set_seed()
-    transforms = utils.build_transforms(hyperparams)
-    loaders = utils.build_loaders(data_dir, transforms, batch_sizes, num_workers, second_stage=(stage == "second"))
     model = utils.build_model(backbone, second_stage=(stage == "second"), num_classes=num_classes, ckpt_pretrained=ckpt_pretrained).cuda()
     model.use_projection_head(False)
     model.eval()
 
-    ## Determine which embeddings to compute
+    ## prep computation
+    transforms = utils.build_transforms(hyperparams)
+    loaders = utils.build_loaders(
+        data_dir, transforms, batch_sizes, num_workers, 
+        second_stage=(stage == "second"), drop_last=False, shuffle_train=False)
     embeddings, labels, rel_paths = [], [], []
 
-    ## val batch size cant be zero
+    ## val set - batch size cant be zero
     embeddings_val, labels_val = utils.compute_embeddings(loaders["valid_loader"], model)
-    if len(embeddings_val) < len(loaders["valid_loader"].dataset.imgs):
-        missed_imgs = len(loaders["valid_loader"].dataset.imgs) - len(embeddings_val)
-        print(f"Warning: missed {missed_imgs} images because batch size was not a multiple of validation dataset size.")
-    rel_paths_val = [item[0][len(root_dir) + 1:] for item in loaders["valid_loader"].dataset.imgs[:len(embeddings_val)]]
+    rel_paths_val = [item[0][len(root_dir) + 1:] for item in loaders["valid_loader"].dataset.imgs]
     embeddings.extend(embeddings_val)
     labels.extend(labels_val)
     rel_paths.extend(rel_paths_val)
 
-    ## train set embeddings
+    ## train set - skipped if zero batch size
     if batch_sizes["train_batch_size"] is not None:
         embeddings_train, labels_train = utils.compute_embeddings(loaders["train_loader"], model)
-        if len(embeddings_train) < len(loaders["train_loader"].dataset.imgs):
-            missed_imgs = len(loaders["train_loader"].dataset.imgs) - len(embeddings_train)
-            print(f"Warning: missed {missed_imgs} images because batch size was not a multiple of training dataset size.")
-        rel_paths_train = [item[0][len(root_dir) + 1:] for item in loaders["train_loader"].dataset.imgs[:len(embeddings_train)]]
+        rel_paths_train = [item[0][len(root_dir) + 1:] for item in loaders["train_loader"].dataset.imgs]
         embeddings.extend(embeddings_train)
         labels.extend(labels_train)
         rel_paths.extend(rel_paths_train)
@@ -118,7 +115,7 @@ def interactive_plots(
 
     ## Reduce dimensionality
     if not perplexity:
-        perplexity = min(100, len(embeddings) // 2)
+        perplexity = min(30, max(5, (len(embeddings) - 1) / 3))
         print(f"tSNE: using a perplexity value of {perplexity}")
     reduced_data, colnames, _ = helpers.embbedings_dimension_reductions(embeddings, perplexity)
 
@@ -127,7 +124,6 @@ def interactive_plots(
     df["paths"] = [os.path.join("..", "..", p) for p in rel_paths]
     df["class"], df["class_str"] = labels, [os.path.basename(os.path.dirname(p)) for p in rel_paths]
     df["dataset"] = df["paths"].apply(lambda x: "validation" if "/val/" in x else "train")
-        
     helpers.bokeh_plot(df, out_path=plot_path, **plot_config)
 
 

bioencoder/scripts/model_explorer.py

@@ -73,7 +73,10 @@ def model_explorer(
     backbone = hyperparams["model"]["backbone"]
     num_classes = hyperparams["model"].get("num_classes", None)
     stage = hyperparams["model"]["stage"]
-
+    img_size = hyperparams.get("img_size", None)
+    if img_size is None:
+        raise ValueError("config must include 'img_size'")
+    
     ## get swa path
     ckpt_pretrained = os.path.join(root_dir, "weights", run_name, stage, "swa")
 
@@ -94,7 +97,7 @@ def model_explorer(
     uploaded_file = st.sidebar.file_uploader("Upload an Image", type=["png", "jpg", "jpeg"])
 
     ## get image transformations
-    transform = utils.get_transforms(hyperparams, valid=False)
+    transform = utils.get_transforms(hyperparams, no_aug=True)
 
     # Load the model and add to cache
     model = load_model(
@@ -105,9 +108,13 @@ def model_explorer(
         )
 
     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)
+        
+        # resize image
+        image_resized = image.resize((img_size, img_size))
 
         # Generate visualizations
         selected = option_menu(None, vis_funcs, icons=['list' for _ in range(len(vis_funcs))], menu_icon="cast", orientation="horizontal")
@@ -123,27 +130,31 @@ def model_explorer(
             layer = st.selectbox("Select a layer", layers.keys())
             module = layers[layer]
             max_acts = st.slider("Max activations", 5, 64, 25)
-            result = vis.visualize_activations(model, module, image, max_acts=max_acts)
+            result = vis.visualize_activations(model, module, image_resized, max_acts=max_acts)
             st.pyplot(result)
 
         elif selected == 'Saliency':
-            result = vis.saliency_map(model, image)
+            result = vis.saliency_map(model, image_resized)
             st.pyplot(result)
 
         elif selected == 'GradCAM':
             # add activation type (Relu, Silu, etc_)
-            layers =[name.split('.')[0] for name, module in model.encoder.named_modules() if isinstance(module, (torch.nn.SiLU, torch.nn.ReLU))]
+            layers =[name.split('.')[0] for name, module in model.encoder.named_modules() \
+                     if isinstance(module, (torch.nn.SiLU, torch.nn.ReLU))]
             layer_set = sorted(set(layers))
             layer = st.selectbox("Select a layer", list(layer_set), index=len(list(layer_set))-1)
-            result = vis.grad_cam(model, model.encoder,image,target_layer=[layer], target_category= None)
+            result = vis.grad_cam(model, model.encoder,image_resized,target_layer=[layer], target_category= None)
             st.pyplot(result)
 
         elif selected == 'ConstrativeCAM':
-            layers =[name.split('.')[0] for name, module in model.encoder.named_modules() if isinstance(module, (torch.nn.SiLU, torch.nn.ReLU))]
+            layers =[name.split('.')[0] for name, module in model.encoder.named_modules() \
+                     if isinstance(module, (torch.nn.SiLU, torch.nn.ReLU))]
             layer_set = sorted(set(layers))
             layer = st.selectbox("Select a layer", list(layer_set), index=len(list(layer_set))-1)
             target = st.selectbox("Select a target", class_names)
-            result = vis.contrast_cam(model, model.encoder,image,target_layer=[layer], target_category=class_names.index(target))
+            result = vis.contrast_cam(
+                model, model.encoder, image_resized,target_layer=[layer], 
+                target_category=class_names.index(target))
             st.pyplot(result)
 
 

bioencoder/scripts/model_explorer_wrapper.py

@@ -15,9 +15,6 @@ def model_explorer_wrapper(config_path):
     process = ["streamlit", "run", script_path , "--", "--config-path", config_path]
     subprocess.run(process)
 
-
-
-    
 def cli():
 
     parser = argparse.ArgumentParser()

bioencoder/vis/helpers.py

@@ -242,7 +242,7 @@ def embbedings_dimension_reductions(data_table, perplexity):
     return np.hstack((pca, tsne)), names, pca_obj
 
 
-def bokeh_plot(df, out_path='plot.html', color_map="viridis", color_classes=None, plot_style=1, 
+def bokeh_plot(df, out_path='plot.html', color_map="jet1", color_classes=None, plot_style=1, 
                point_size=10, **kwargs):
     """
     Plot a scatter plot of the PCA and t-SNE dimensions of the data using bokeh.
@@ -264,8 +264,9 @@ class labels of the images).
         raise ValueError("The dataframe must have columns 'paths' and 'class'")      
 
     unique_classes = df['class'].unique()
- 
-   
+    unique_datasets = df['dataset'].unique()
+    markers = ['circle', 'square']  # Define markers for each group
+
     ## Color management
     if color_classes:
         assert len(unique_classes) == len(color_classes), (
@@ -284,10 +285,12 @@ class labels of the images).
         colors_str = ['#%02x%02x%02x' % tuple(c[:3]) for c in colors_raw]
         df['color'] = colors_str
 
+        
     source = ColumnDataSource(df)
     bplot.output_file(out_path)
 
     if plot_style == 1:
+        div = Div(text="")
         tooltip = """
         <div>
             <div>
@@ -306,18 +309,12 @@ class labels of the images).
         hover1 = HoverTool(tooltips=tooltip)
         tools0 = [t() for t in TOOLS] + [hover0]
         tools1 = [t() for t in TOOLS] + [hover1]
-        pca = bplot.figure(tools=tools0)
-        pca.scatter('PC1', 'PC2', color='color', source=source, size=point_size)
-        tsne = bplot.figure(tools=tools1)
-        tsne.scatter('tSNE-0', 'tSNE-1', color='color', source=source, size=point_size)
-        p = bplot.gridplot([[pca, tsne]])
-        bplot.show(p)
-        
+
     elif plot_style == 2:
         div = Div(text="")
         hover=HoverTool(
                 tooltips = [
-                ("class_str", "@class_str"),
+                ("Class", "@class_str"),
                 ]
         )
         hover.callback = CustomJS(args=dict(div=div, ds=source), code="""
@@ -333,11 +330,28 @@ class labels of the images).
             """)
         tools0 = [t() for t in TOOLS] + [hover]
         tools1 = [t() for t in TOOLS] + [hover]
-        pca = bplot.figure(tools=tools0)
-        pca.scatter('PC1', 'PC2', color='color', source=source, size=point_size)
-        tsne = bplot.figure(tools=tools1)
-        tsne.scatter('tSNE-0', 'tSNE-1', color='color', source=source, size=point_size)
-        p = bplot.gridplot([[pca, tsne]])
-        show(layouts.row(p, div))
+        
+    # Create figures
+    pca = bplot.figure(tools=tools0, title="PCA", match_aspect=True)
+    tsne = bplot.figure(tools=tools1, title="t-SNE", match_aspect=True)
+    
+    # Store renderers for dataset legend
+    legend_items_dataset = []
+    
+    # Scatter plots with different markers for datasets
+    for dataset, marker in zip(unique_datasets, markers):
+        dataset_source = ColumnDataSource(df[df['dataset'].astype(str) == dataset])  # Filter dataset-specific data
+        r = pca.scatter('PC1', 'PC2', source=dataset_source, color='color', size=point_size, marker=marker)
+        tsne.scatter('tSNE-0', 'tSNE-1', source=dataset_source, color='color', size=point_size, marker=marker)
+        legend_items_dataset.append(LegendItem(label=str(dataset), renderers=[r]))
+    
+    # Create and add horizontal legend for dataset markers
+    legend_dataset = Legend(items=legend_items_dataset, orientation="horizontal")
+    pca.add_layout(legend_dataset, 'below')
+    
+    # Display plots
+    p = bplot.gridplot([[pca, tsne]])
+    show(layouts.row(p, div))
+
 
     return p

bioencoder/vis/methods.py

@@ -110,7 +110,8 @@ def saliency_map(model, img, device = 'cuda', save_path = None):
         plt.savefig(save_path)
     return fig
 
-def grad_cam(model, module, img, target_layer = ["4"], target_category= None, device = 'cuda', save_path = None):
+def grad_cam(model, module, img, target_layer = ["4"], 
+             target_category= None, device = 'cuda', save_path = None):
 
     for param in model.parameters():
         param.requires_grad = True

bioencoder_configs/plot_stage1.yml

@@ -4,21 +4,19 @@ model:
   stage: first # Training stage: 'first' for initial training, 'second' for fine-tuning
 
 dataloaders:
-  train_batch_size: 2 # Batch size for training data; ensure validation set size is a multiple of this number
-  valid_batch_size: 2 # Batch size for validation data
-  num_workers: 32 # Number of CPU threads for data loading; should not exceed the number of CPU cores
+  train_batch_size: 20 # Larger is faster; no value or removing this line will not include training data
+  valid_batch_size: 20 # Larger is faster; val data is always plotted
+  num_workers: 32 # Should not exceed available CPU cores
 
-img_size: 384 # Image size for training and validation
+img_size: 384 # image size used for training
 
-plot_style: 1 # (1: pictogram above point, 2: pictogram next to plot panel)
-
-color_classes: # overrides color_map
-  #class1: "#FFD467"
-  #class2: "#4DC9F2"
+perplexity: 30 # for tSNE<; cannot be larger than dataset
 
+plot_style: 2 # (1: pictogram above point, 2: pictogram next to plot panel)
 point_size: 10  ## size of points in scatter plot
-
 color_map: 'Set1' # Default color map; see https://matplotlib.org/stable/users/explain/colors/colormaps.html for options
-
+#color_classes: # overrides color_map
+  #class1: "#FFD467"
+  #class2: "#4DC9F2"