from torchvision import models
import torch.nn as nn

# Class names (IMPORTANT for prediction later)
CLASS_NAMES = ['cardboard', 'glass', 'metal', 'paper', 'plastic', 'trash']


def get_model():
    """
    Returns the ResNet50 model with modified final layer
    """
    model = models.resnet50(pretrained=False)  # DO NOT change this

    # Replace final layer for 6 classes
    model.fc = nn.Linear(model.fc.in_features, len(CLASS_NAMES))

    return model