models.py

import torch
from torchvision.models import  ViT_B_16_Weights, EfficientNet_B2_Weights


def get_vit_16_base_transformer():
    vit_b_16_model = torch.load("models/ViT_16_base_101_classes_pretrained_custom_head.pth", map_location = torch.device("cpu"))
    vit_b_16_transforms = ViT_B_16_Weights.DEFAULT.transforms()

    return vit_b_16_model, vit_b_16_transforms

def get_effnet_b2():
    eff_net_b2_model = torch.load("models/effnet_b2_101_classes_pretrained_custom_head.pth", map_location = torch.device("cpu"))
    eff_net_b2_transforms = EfficientNet_B2_Weights.DEFAULT.transforms()

    return eff_net_b2_model, eff_net_b2_transforms