diff --git a/examples/Cats_vs_Dogs/cat_dog_train.yml b/examples/Cats_vs_Dogs/cat_dog_train.yml index 10a98a2..b8302fb 100644 --- a/examples/Cats_vs_Dogs/cat_dog_train.yml +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index d45f40b..590daa1 100644 --- a/pyproject.toml +++ b/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" diff --git a/src/animl/__init__.py b/src/animl/__init__.py index 6e22f67..5907d2d 100644 --- a/src/animl/__init__.py +++ b/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', diff --git a/src/animl/classification.py b/src/animl/classification.py index e8533fb..6e8932e 100644 --- a/src/animl/classification.py +++ b/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 diff --git a/src/animl/config/train.yml b/src/animl/config/train.yml index cdab98c..a3ef650 100644 --- a/src/animl/config/train.yml +++ b/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 diff --git a/src/animl/generator.py b/src/animl/generator.py index d27148e..be14641 100755 --- a/src/animl/generator.py +++ b/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, diff --git a/src/animl/model_architecture.py b/src/animl/model_architecture.py index 474b02b..87c44b5 100644 --- a/src/animl/model_architecture.py +++ b/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 diff --git a/src/animl/test.py b/src/animl/test.py index c522f32..17e8350 100644 --- a/src/animl/test.py +++ b/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 diff --git a/src/animl/train.py b/src/animl/train.py index 3511acd..2d5f74d 100755 --- a/src/animl/train.py +++ b/src/animl/train.py @@ -53,11 +53,21 @@ def save_classifier(model, None ''' Path(out_dir).mkdir(parents=True, exist_ok=True) - - # get model parameters and add to stats - checkpoint = {'model': model.state_dict(), - 'stats': stats} - # save optimizer and scheduler state dicts if they are provided + if model.__class__.__name__=="BioClip": + # save only parameters that are changeable: lora and classifier + trainable_state_dict = { + k: v for k, v in model.state_dict().items() + if "lora" in k or "classifier" in k + } + checkpoint = { + 'model': trainable_state_dict, + 'stats': stats + } + else: + # get model parameters and add to stats + checkpoint = {'model': model.state_dict(), + 'stats': stats} + # save optimizer, scheduler, and scaler state dicts if they are provided if optimizer is not None or scheduler is not None: checkpoint['epoch'] = epoch if optimizer is not None: @@ -99,7 +109,7 @@ def load_classifier_checkpoint(model_path, model, optimizer, scheduler, scaler, # load state dict and apply weights to model print(f'Resuming from epoch {start_epoch}') checkpoint = torch.load(open(f'{model_path}/{start_epoch}.pt', 'rb'), map_location=device) - model.load_state_dict(checkpoint['model']) + model.load_state_dict(checkpoint['model'], strict=False) # Model is assumed to be on the correct device already (moved in main before optimizer creation) # load optimzier state if available @@ -118,13 +128,13 @@ def load_classifier_checkpoint(model_path, model, optimizer, scheduler, scaler, if 'scaler' in checkpoint and scaler is not None: scaler.load_state_dict(checkpoint['scaler']) - # get last epoch from model if avialble + # get last epoch from model if available if 'epoch' in checkpoint: return checkpoint['epoch'] else: return start_epoch else: - # no save state found; stasrt anew + # no save state found; start anew print('No model state found, starting new model') return 0 @@ -340,6 +350,7 @@ def train_classifier(cfg): file_col = cfg.get('file_col', 'filepath') label_col = cfg.get('label_col', 'species') resize_width, resize_height = cfg.get('image_size', [480,480]) + architecture=cfg['architecture'] # check if GPU is available device = cfg.get('device', 'cpu') @@ -371,14 +382,16 @@ def train_classifier(cfg): file_col=file_col, label_col=label_col, crop=crop, augment=cfg.get('augment', True), resize_height=resize_height, resize_width=resize_width, - cache_dir=cfg.get('cache_folder', None)) + cache_dir=cfg.get('cache_folder', None), + architecture=architecture) dl_val = train_dataloader(validate_dataset, categories, batch_size=cfg.get('val_batch_size', 16), num_workers=cfg.get('num_workers', NUM_THREADS), file_col=file_col, label_col=label_col, 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) # set up model optimizer if cfg.get("optimizer", "AdamW") == 'AdamW': @@ -428,8 +441,10 @@ def train_classifier(cfg): print(f"Using learning rate : {scheduler.get_last_lr()[0]}") if current_epoch > frozen_epochs: - for param in model.parameters(): - param.requires_grad = True + for name, param in model.named_parameters(): + if architecture != "bioclip_2" or "lora" in name: + #for bioclip, we only want to unfreeze the lora parameters + param.requires_grad = True loss_train, oa_train = _train_classifier_helper(dl_train, model, optim, scheduler, scaler=scaler, device=device, mixed_precision=mixed_precision, progress=progress)