removed py39 req / fixed small bugs

Author: mluerig

bioencoder/core/utils.py

@@ -324,7 +324,7 @@ def build_optim(model, optimizer_params, scheduler_params, loss_params):
     return {"criterion": criterion, "optimizer": optimizer, "scheduler": scheduler, "loss_optimizer": loss_optimizer}
 
 
-def compute_embeddings(loader, model):
+def compute_embeddings(loader, model, scaler=None):
     """Computes the embeddings and corresponding labels for a dataset.
 
     Parameters:
@@ -342,7 +342,11 @@ def compute_embeddings(loader, model):
 
     for images, labels in loader:
         images = images.cuda()
-        embed = model(images)
+        if scaler:
+            with torch.cuda.amp.autocast():
+                embed = model(images)
+        else:
+            embed = model(images)
         if total_embeddings is None:
             total_embeddings = embed.detach().cpu()
             total_labels = labels.detach().cpu()
@@ -644,7 +648,8 @@ def save_augmented_sample(data_dir, transform, n_samples, seed):
     # Load dataset
     dataset = ImageFolder(root=os.path.join(data_dir, "train"))
     save_dir = os.path.join(data_dir, "aug_sample")
-    shutil.rmtree(save_dir)
+    if os.path.isdir(save_dir):
+        shutil.rmtree(save_dir)
     os.makedirs(save_dir, exist_ok=True)
 
     ## reverse image net transforms

bioencoder/scripts/model_explorer.py

@@ -71,9 +71,8 @@ def model_explorer(
 
     ## parse config
     backbone = hyperparams["model"]["backbone"]
-    num_classes = hyperparams["model"]["num_classes"]
+    num_classes = hyperparams["model"].get("num_classes", None)
     stage = hyperparams["model"]["stage"]
-    device = hyperparams.get("device", "cuda")
 
     ## get swa path
     ckpt_pretrained = os.path.join(root_dir, "weights", run_name, stage, "swa")
@@ -106,18 +105,10 @@ def model_explorer(
         )
 
     if uploaded_file is not None:
-        
-        ## load image
-        image = Image.open(uploaded_file)
-        
         # Display the uploaded image
+        image = Image.open(uploaded_file).convert('RGB')
         st.sidebar.image(image, caption="Input Image", use_column_width=True)
 
-        ## apply transforms and load into device
-        image = np.asarray(image)
-        image = transform(image=image)["image"]
-        image = image.unsqueeze(0).to(device)
-    
         # Generate visualizations
         selected = option_menu(None, vis_funcs, icons=['list' for _ in range(len(vis_funcs))], menu_icon="cast", orientation="horizontal")
         if selected == 'Filters':
@@ -154,7 +145,6 @@ def model_explorer(
             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))
             st.pyplot(result)
-        
 
 
 if __name__ == "__main__":

bioencoder/vis/methods.py

@@ -52,7 +52,7 @@ def visualize_activations(model, module, img, max_acts = 64, save_path = None, d
         Plots the activations of a module recorded during a forward pass on an image
     """
     model.to(device)
-    # img_t = preprocess_image(img).to(device)
+    img_t = preprocess_image(img).to(device)
     acts = [0]
 
     def hook_fn(self, input, output):
@@ -87,7 +87,7 @@ def saliency_map(model, img, device = 'cuda', save_path = None):
 
     model.eval()
     model.to(device)    
-    # img_t = preprocess_image(img).to(device)
+    img_t = preprocess_image(img).to(device)
     img_t.requires_grad = True
     img_t.retain_grad() #added this line 
 
@@ -120,7 +120,7 @@ def grad_cam(model, module, img, target_layer = ["4"], target_category= None, de
     grad_cam = GradCam(model = model, feature_module = module,
                     target_layer_names = target_layer, use_cuda = use_cuda)
 
-    # img_t = preprocess_image(img).to(device)
+    img_t = preprocess_image(img).to(device)
 
     grayscale_cam = grad_cam(img_t, target_category)
 
@@ -167,7 +167,7 @@ def contrast_cam(model, module, img, target_layer = ["4"], target_category= None
     contrast_cam = ContrastCam(model = model, feature_module = module,
                     target_layer_names = target_layer, use_cuda = use_cuda)
 
-    # img_t = preprocess_image(img).to(device)
+    img_t = preprocess_image(img).to(device)
 
     assert(target_category != None), "Please specify a target category"
     grayscale_cam = contrast_cam(img_t, target_category)
@@ -204,4 +204,4 @@ def contrast_cam(model, module, img, target_layer = ["4"], target_category= None
 
     if save_path:
         fig1.savefig(save_path)
-    return fig1
+    return fig1

config-templates/train_stage1.yml

@@ -31,14 +31,17 @@ criterion:
   params:
     temperature: 0.1 # Temperature parameter for contrastive loss
 
-img_size: 384 # Image size for training and validation
+img_size: &size 384  # Image size for training and validation
 
 augmentations: # augmentations to be applied - see https://github.com/agporto/BioEncoder/blob/main/help/05-options.md#augmentations
   sample_save: True # Whether to save a sample of augmented images
   sample_n: 10 # Number of augmented image samples per class to save
   sample_seed: 42 # Seed for random sample
   transforms: 
     - RandomResizedCrop: # Randomly resize and crop the image
+        height: *size
+        width: *size
+        scale:  !!python/tuple [0.7,1]
     - Flip: # Randomly flip the image horizontally
     - RandomRotate90: # Randomly rotate the image by 90 degrees
     - MedianBlur: # Apply median blur with a probability

config-templates/train_stage2.yml

@@ -32,14 +32,17 @@ criterion:
     classes: 100 # Number of classes (adjust based on actual number of classes)
     smoothing: 0.01 # Smoothing factor for label smoothing
 
-img_size: 384 # Image size for training and validation
+img_size: &size 384  # Image size for training and validation
 
 augmentations: # augmentations to be applied - see https://github.com/agporto/BioEncoder/blob/main/help/05-options.md#augmentations
   sample_save: True # Whether to save a sample of augmented images
   sample_n: 10 # Number of augmented image samples per class to save
   sample_seed: 42 # Seed for random sample
-  transforms: 
+  transforms:
     - RandomResizedCrop: # Randomly resize and crop the image
+        height: *size
+        width: *size
+        scale:  !!python/tuple [0.7,1]
     - Flip: # Randomly flip the image horizontally
     - RandomRotate90: # Randomly rotate the image by 90 degrees
     - MedianBlur: # Apply median blur with a probability

pyproject.toml

@@ -10,7 +10,7 @@ authors = [
 	]
 description = "A metric learning toolkit"
 readme = "README.md"
-requires-python = "==3.9.*"
+requires-python = ">=3.9"
 keywords = ["metric learning", "biology"]
 dynamic = ["dependencies"]
 version = "0.3.1"