import torch
import torchvision
from torchvision import models
from torch import nn

device = 'cpu'
def create_effnetb2(num_output_classes: int):
    """
    Returns an effnetb2 feature extractor model with all layers except classifier layer frozen and the corresponding transforms for data preprocessing
    
    Args
        num_output_classes (int) : The number of classes in the classifier head
    """
    effnet_b2_weights = models.EfficientNet_B2_Weights.DEFAULT
    effnet_b2 = models.efficientnet_b2(weights= effnet_b2_weights).to(device)
    effnet_transforms = effnet_b2_weights.transforms()
    for params in effnet_b2.parameters():
        params.requires_grad = False
    
    effnet_b2.classifier = nn.Sequential(
    nn.Dropout(p= 0.3, inplace= True),
    nn.Linear(in_features= 1408, out_features= num_output_classes, bias= True)
    ).to(device)
    
    return effnet_b2, effnet_transforms