319 add bioclip training

examples/Cats_vs_Dogs/cat_dog_train.yml

@@ -21,7 +21,7 @@ class_list_index: 'id'
 class_list_label: 'class'
 
 # training hyperparameters
-architecture: "efficientnet_v2_m" # or choose "convnext_base"
+architecture: "efficientnet_v2_m" # or choose "convnext_base" or "bioclip_2"
 image_size: [480, 480] # [width, height] preprocess images to this size (architecture dependent. some models require a certain size)
 crop: False # train on cropped images or not
 batch_size: 12

pyproject.toml

@@ -12,16 +12,21 @@ authors = [{name = "Kyra Swanson", email = "tswanson@sdzwa.org" }]
 requires-python = ">=3.12"
 dependencies = [
     "dill>=0.4.0",
+    "huggingface_hub>=0.23.0",
     "numpy>=2.0.2", 
     "onnxruntime-gpu>=1.19.2", 
     "pandas>=2.2.2,<3.0.0", 
     "pillow>=11.0.0", 
     "pyexiftool>=0.5.6",
-    "opencv-python>=4.12.0.88", 
+    "opencv-python>=4.12.0.88",
+    "open_clip_torch>=2.24.0",
+    "peft>=0.12.0",
+    "safetensors>=0.4.3",
     "scikit-learn>=1.5.2", 
     "timm>=1.0.9",
     "torch>=2.6.0",
     "torchvision>=0.21.0",
+    "transformers>=4.40.0",
     "tqdm>=4.66.5",
     "ultralytics>=8.3.95",
     "wget>=3.2"

src/animl/__init__.py

@@ -33,8 +33,8 @@
                                    load_json, load_yaml, save_data, save_json,
                                    save_yaml, sequence_calculation,)
 from animl.generator import (Letterbox, ManifestGenerator, TrainGenerator,
-                             collate_fn, image_to_tensor, manifest_dataloader,
-                             train_dataloader,)
+                             collate_fn, image_to_tensor,
+                             manifest_dataloader, train_dataloader,)
 from animl.model_architecture import (ConvNeXtBase, EfficientNet, MD_LABELS,
                                       MD_MODELS, MEGADETECTORv5_SIZE,
                                       MEGADETECTORv5_STRIDE,
@@ -98,8 +98,8 @@
            'export_train_val_test', 'export_yolo', 'extract_frames',
            'extract_miew_embeddings', 'file_management', 'from_config',
            'from_paths', 'general', 'generator', 'get_animals', 'get_empty',
-           'get_frame_as_image', 'get_iou', 'get_onnx_device', 'get_torch_device',
-           'get_version', 'image_to_tensor', 'inference', 'init_seed',
+           'get_frame_as_image', 'get_iou', 'get_onnx_device',
+           'get_torch_device', 'get_version', 'image_to_tensor', 'inference', 'init_seed',
            'l2_norm', 'letterbox', 'list_models', 'load_class_list',
            'load_classifier', 'load_classifier_checkpoint', 'load_data',
            'load_detector', 'load_json', 'load_miew', 'load_yaml',

src/animl/classification.py

@@ -15,7 +15,7 @@
 import onnxruntime
 
 from animl import generator, file_management
-from animl.model_architecture import EfficientNet, ConvNeXtBase
+from animl.model_architecture import EfficientNet, ConvNeXtBase, BioClip
 from animl.utils.general import (get_torch_device, get_onnx_device, softmax,
                                  tensor_to_onnx, NUM_THREADS)
 
@@ -56,7 +56,7 @@ def load_classifier(model_path: str,
 
     # Create a new model instance for training (pytorch only)
     if model_path.is_dir():
-        supported_architectures = ["efficientnet_v2_m", "convnext_base", ]
+        supported_architectures = ["efficientnet_v2_m", "convnext_base", "bioclip_2"]
         if architecture not in supported_architectures:
             raise ValueError(f"""Unsupported architecture: {architecture}.
                              Supported architectures are: {supported_architectures}""")
@@ -69,6 +69,8 @@ def load_classifier(model_path: str,
             model = EfficientNet(num_classes, device=device)
         elif architecture == "convnext_base":
             model = ConvNeXtBase(num_classes)
+        elif architecture == "bioclip_2":
+            model = BioClip(num_classes)
         else:  # can only resume models from a directory at this time
             raise AssertionError('Please provide the correct model')
         return model, class_list, start_epoch
@@ -79,7 +81,7 @@ def load_classifier(model_path: str,
         start_time = time()
         # PyTorch dict
         if model_path.suffix == '.pt':
-            supported_architectures = ["efficientnet_v2_m", "convnext_base"]
+            supported_architectures = ["efficientnet_v2_m", "convnext_base", "bioclip_2"]
             if architecture not in supported_architectures:
                 raise ValueError(f"""Unsupported architecture: {architecture}.
                                  Supported architectures are: {supported_architectures}""")
@@ -100,6 +102,13 @@ def load_classifier(model_path: str,
                 model.to(device)
                 model.eval()
                 model.framework = "ConvNeXt-Base"
+            elif architecture == "bioclip_2":
+                model = BioClip(num_classes, tune=False)
+                checkpoint = torch.load(model_path, map_location=device)
+                model.load_state_dict(checkpoint['model'], strict=False)
+                model.to(device)
+                model.eval()
+                model.framework = "BioClip"
         # PyTorch full modelspeak
         elif model_path.suffix == '.pth':
             # check to make sure GPU is available if chosen

src/animl/config/train.yml

@@ -24,7 +24,7 @@ class_list_index: 'id'
 class_list_label: 'class'
 
 # training hyperparameters
-architecture: "efficientnet_v2_m" # or choose "convnext_base"
+architecture: "efficientnet_v2_m" # or choose "convnext_base" or "bioclip_2"
 image_size: [480, 480] # [width, height] preprocess images to this size (architecture dependent. some models require a certain size)
 crop: True # train on cropped images or not
 crop_coord: 'relative' #relative or absolute bounding box coordinates

src/animl/generator.py

@@ -14,9 +14,10 @@
 import torch
 from torch import Tensor
 from torch.utils.data import Dataset, DataLoader
+from torchvision.transforms.functional import InterpolationMode
 from torchvision.transforms.v2 import (Compose, Resize, ToImage, ToDtype, Pad, RandomHorizontalFlip,
                                        RandomAffine, RandomGrayscale, RandomApply,
-                                       ColorJitter, GaussianBlur)
+                                       ColorJitter, GaussianBlur, Normalize)
 
 
 from animl.model_architecture import SDZWA_CLASSIFIER_SIZE
@@ -112,6 +113,41 @@ def image_to_tensor(file_path, letterbox, resize_width, resize_height):
     return img_tensor, [file_path], [frame], torch.tensor([(height, width)])
 
 
+def _get_model_transforms(resize_height, resize_width, architecture=None):
+    """
+    Generates the image preprocessing pipeline matching the model architecture.
+
+    Args:
+        resize_height (int): resize image height
+        resize_width (int): resize image width
+        architecture (str, optional): Model architecture name. If None, falls
+            back to the default transform pipeline.
+
+    Returns:
+        torchvision.transforms.v2.Compose: Complete preprocessing pipeline.
+    """
+    if architecture == "bioclip_2":
+        if resize_width != 224 or resize_height != 224:
+            resize_width, resize_height = (224,224)
+            print(
+                "[WARNING] Changing resize_width and resize_height to 224x224 "
+                "to satisfy BioClip input requirements."
+            )
+
+        return Compose([Letterbox(resize_height = resize_height,
+                                  resize_width = resize_width,
+                                  color = 127,
+                                  interpolation_mode = InterpolationMode.BICUBIC),
+                        ToImage(),
+                        ToDtype(torch.float32, scale=True),
+                        Normalize(mean=[0.48145466, 0.4578275, 0.40821073],
+                                  std=[0.26862954, 0.26130258, 0.27577711]),])
+    # default transformations
+    return Compose([Resize((resize_height, resize_width)),
+                    ToImage(),
+                    ToDtype(torch.float32, scale=True),])
+
+
 class ManifestGenerator(Dataset):
     '''
     Data generator that crops images on the fly, requires relative bbox coordinates,
@@ -291,6 +327,7 @@ class TrainGenerator(Dataset):
         - resize_height: size in pixels for input height
         - resize_width: size in pixels for input width
         - cache_dir: if not None, use given cache directory to store preprocessed images
+        - architecture: model architecture
     '''
     def __init__(self, x: pd.DataFrame,
                  classes: dict,
@@ -301,7 +338,8 @@ def __init__(self, x: pd.DataFrame,
                  crop: bool = True,
                  crop_coord: str = 'relative',
                  augment: bool = False,
-                 cache_dir: str = None):
+                 cache_dir: str = None,
+                 architecture: str = None):
         self.x = x.reset_index(drop=True)
         self.resize_height = int(resize_height)
         self.resize_width = int(resize_width)
@@ -336,16 +374,10 @@ def __init__(self, x: pd.DataFrame,
                                 # adjust brightness and contrast for varying lighting conditions
                                 ColorJitter(brightness=0.2, contrast=0.2)
                                 ])
+        self.transform = _get_model_transforms(resize_height, resize_width, architecture)
         if self.augment:
             print("Applying augmentations")
-            self.transform = Compose([augmentations,  # augmentations
-                                      Resize((self.resize_height, self.resize_width)),
-                                      ToImage(),
-                                      ToDtype(torch.float32, scale=True),])
-        else:
-            self.transform = Compose([Resize((self.resize_height, self.resize_width)),
-                                      ToImage(),
-                                      ToDtype(torch.float32, scale=True),])
+            self.transform = Compose(augmentations.transforms + self.transform.transforms)
         self.categories = {c: idx for idx, c in classes.items()}
 
     def __len__(self):
@@ -442,7 +474,8 @@ def train_dataloader(manifest: pd.DataFrame,
                      augment: bool = False,
                      batch_size: int = 1,
                      num_workers: int = 1,
-                     cache_dir: str = None):
+                     cache_dir: str = None,
+                     architecture: str = None):
     '''
     Loads a dataset for training and wraps it in a PyTorch DataLoader object.
 
@@ -461,6 +494,7 @@ def train_dataloader(manifest: pd.DataFrame,
         batch_size (int): size of each batch
         num_workers (int): number of processes to handle the data
         cache_dir (str): if not None, use given cache directory
+        architecture (str): determines transforms used. if None, uses default transforms
 
     Returns:
         dataloader object
@@ -474,7 +508,8 @@ def train_dataloader(manifest: pd.DataFrame,
                                       resize_height=resize_height,
                                       resize_width=resize_width,
                                       augment=augment,
-                                      cache_dir=cache_dir)
+                                      cache_dir=cache_dir,
+                                      architecture=architecture)
 
     dataLoader = DataLoader(dataset=dataset_instance,
                             batch_size=batch_size,

src/animl/model_architecture.py

@@ -3,11 +3,12 @@
 
 @ Kyra Swanson 2023
 """
+import open_clip
+from peft import LoraConfig, get_peft_model
 import torch
 import torch.nn as nn
 from torchvision.models import efficientnet, convnext_base, ConvNeXt_Base_Weights
 
-
 MEGADETECTORv5_SIZE = 1280
 MEGADETECTORv5_STRIDE = 64
 MD_LABELS = {0: "empty", 1: "animal", 2: "human",  3: "vehicle"}
@@ -71,3 +72,47 @@ def forward(self, x):
         Forward pass (prediction).
         '''
         return self.model(x)
+
+
+class BioClip(nn.Module):
+    '''
+    Construct the BioClip2 model architecture.
+    '''
+    def __init__(self, num_classes, tune=False):
+        super(BioClip,self).__init__()
+        # load the BioClip2 vision encoder pre-trained on TreeOfLife-200M
+        full_model, self.preprocess_train, self.preprocess_val = (
+            open_clip.create_model_and_transforms('hf-hub:imageomics/bioclip-2')
+        )
+        self.model = full_model.visual
+        embedding_dim = self.model.output_dim
+        # set up low-rank adaptation
+        config = LoraConfig(
+                r=16,
+                lora_alpha=32,
+                target_modules=["attn","c_fc","c_proj"],
+                lora_dropout=0.0,
+                bias="none",
+                modules_to_save=None
+        )
+        # this freezes the base visual encoder and injects trainable LoRA layers
+        self.model = get_peft_model(self.model, config)
+
+        if not tune:
+            for param in self.model.parameters():
+                param.requires_grad = False
+
+        # Add a classifier layer
+        self.classifier = nn.Linear(in_features=embedding_dim, out_features=num_classes)
+
+    def forward(self,x):
+        '''
+        Forward pass (prediction)
+        '''
+        # only use the visual encoder of the BioClip2 model
+        features = self.model(x)
+        # features.size(): [B x 768]
+        # normalize so the model learns based off direction only
+        features = features / features.norm(dim=-1, keepdim=True)
+        logits = self.classifier(features)
+        return logits

src/animl/test.py

@@ -78,6 +78,7 @@ def test_classifier(cfg):
 
     crop = cfg.get('crop', False)
     resize_width, resize_height = cfg.get('image_size', [480,480])
+    architecture=cfg['architecture']
 
     # check if GPU is available
     device = cfg.get('device', 'cpu')
@@ -89,7 +90,7 @@ def test_classifier(cfg):
     model, classes = load_classifier(cfg['active_model'],
                                      cfg['class_file'],
                                      device=device,
-                                     architecture=cfg['architecture'])
+                                     architecture=architecture)
 
     class_list_label = cfg.get('class_list_label', 'class')
     class_list_index = cfg.get('class_list_index', 'id')
@@ -105,7 +106,8 @@ def test_classifier(cfg):
                                label_col=cfg.get('label_col', 'species'),
                                crop=crop, augment=False,
                                resize_height=resize_height, resize_width=resize_width,
-                               cache_dir=cfg.get('cache_folder', None))
+                               cache_dir=cfg.get('cache_folder', None),
+                               architecture=architecture)
     # get predictions
     pred, true, paths = _test_classifer_helper(dl_test, model, device)
     # calculate precision and recall