Spaces:
Sleeping
Sleeping
| import torch | |
| from model.BiSeNet.build_bisenet import BiSeNet | |
| from model.BiSeNetV2.model import BiSeNetV2 | |
| # Model loading function | |
| def loadModel(model:str = 'bisenet', device: str = 'cpu', weights:str='weight_Base.pth') -> BiSeNet | BiSeNetV2: | |
| """ | |
| Load the BiSeNet or BiSeNetV2 model and move it to the specified device. | |
| This function supports loading different versions of the model based on the provided `model` argument. | |
| The model weights are loaded from the specified `weights` file. | |
| The model is set to evaluation mode after loading. | |
| Args: | |
| model (str): The type of model to load. Options are 'bisenet', 'bisenet_base', 'bisenet_best', 'bisenetv2', 'bisenetv2_base', 'bisenetv2_best'. | |
| Default is 'bisenet'. | |
| device (str): Device to load the model onto ('cpu' or 'cuda'). Default is 'cpu'. | |
| weights (str): weights file to be loaded. Default is 'weight_Base.pth'. | |
| Returns: | |
| model (BiSeNet | BiSeNetV2): The loaded BiSeNet or BiSeNetV2 model. | |
| """ | |
| match model.lower() if isinstance(model, str) else model: | |
| case 'bisenet' | 'bisenet_base' | 'bisenet_best': | |
| model = BiSeNet(num_classes=19, context_path='resnet18').to(device) | |
| modelStateDict = torch.load(f'./weights/BiSeNet/{weights}', map_location=device) | |
| model.load_state_dict(modelStateDict['model_state_dict'] if 'model_state_dict' in modelStateDict else modelStateDict) | |
| case 'bisenetv2' | 'bisenetv2_base' | 'bisenetv2_best': | |
| model = BiSeNetV2(n_classes=19).to(device) | |
| modelStateDict = torch.load(f'./weights/BiSeNetV2/{weights}', map_location=device) | |
| model.load_state_dict(modelStateDict['model_state_dict'] if 'model_state_dict' in modelStateDict else modelStateDict) | |
| case _: raise NotImplementedError(f"Model {model} is not implemented. Please choose 'bisenet' or 'bisenetv2'.") | |
| model.eval() | |
| return model |