Imageomics / pybioclip

Python package that simplifies using the BioCLIP foundation model.
MIT License
12 stars 3 forks source link

Using Other CLIP Checkpoints #23

Closed work4cs closed 1 month ago

work4cs commented 1 month ago

We can give An Example in READ.ME for using other openclip style checkpoints:

from bioclip import CustomLabelsClassifier

classifier = CustomLabelsClassifier(['Schistocerca americana', 'Schistocerca lineata'], device='cuda', model_str='hf-hub:laion/CLIP-ViT-B-16-laion2B-s34B-b88K')
predictions = classifier.predict("image.png")

Also could even allow openai checkpoint or local checkpoint by passing pretrained parameters.

class BaseClassifier(object):
    def __init__(self, device: Union[str, torch.device] = 'cpu', model_str: str = MODEL_STR, pretrained: str = 'openai'):
        self.model = create_bioclip_model(device=device, model_str=model_str, pretrained: str =pretrained )
def create_bioclip_model(model_str, device="cuda", pretrained: str = 'openai'):
    model = create_model(model_str, output_dict=True, require_pretrained=True, pretrained: str=pretrained)