| # YOLOP by hustvl, MIT License | |
| dependencies = ['torch'] | |
| import torch | |
| from lib.utils.utils import select_device | |
| from lib.config import cfg | |
| from lib.models import get_net | |
| from pathlib import Path | |
| import os | |
| def yolop(pretrained=True, device="cpu"): | |
| """Creates YOLOP model | |
| Arguments: | |
| pretrained (bool): load pretrained weights into the model | |
| wieghts (int): the url of pretrained weights | |
| device (str): cuda device i.e. 0 or 0,1,2,3 or cpu | |
| Returns: | |
| YOLOP pytorch model | |
| """ | |
| device = select_device(device = device) | |
| model = get_net(cfg) | |
| if pretrained: | |
| path = os.path.join(Path(__file__).resolve().parent, "weights/End-to-end.pth") | |
| checkpoint = torch.load(path, map_location= device) | |
| model.load_state_dict(checkpoint['state_dict']) | |
| model = model.to(device) | |
| return model | |