model/modelLoading.py

import torch

from model.BiSeNet.build_bisenet import BiSeNet


# %% load model

def loadModel(model:str = 'bisenet', device: str = 'cpu')->BiSeNet:
    """
    Load the specified model and move it to the given device.
    
    Args:
        model (str): model to be loaded.
        device (str): Device to load the model onto ('cpu' or 'cuda').
        
    Returns:
        model (BiSeNet): The loaded BiSeNet model.
    """
    match model.lower() if isinstance(model, str) else model:
        case 'bisenet': model = loadBiSeNet(device)
        case _: raise NotImplementedError(f"Model {model} is not implemented. Please choose 'bisenet' .")
    
    return model


# BiSeNet model loading function
def loadBiSeNet(device: str = 'cpu') -> BiSeNet:
    """
    Load the BiSeNet model and move it to the specified device.
    
    Args:
        device (str): Device to load the model onto ('cpu' or 'cuda').
        
    Returns:
        model (BiSeNet): The loaded BiSeNet model.
    """
    model = BiSeNet(num_classes=19, context_path='resnet18').to(device)
    model.load_state_dict(torch.load('./weights/BiSeNet/weightADV.pth', map_location=device)['model_state_dict'])
    model.eval()
    
    return model