Spaces:
Sleeping
Sleeping
File size: 1,216 Bytes
5a17bb3 c06c582 5a17bb3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 | 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 |