You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

101 lines
3.4KB

  1. """File for accessing YOLOv5 via PyTorch Hub https://pytorch.org/hub/
  2. Usage:
  3. import torch
  4. model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True, channels=3, classes=80)
  5. """
  6. dependencies = ['torch', 'yaml']
  7. import os
  8. import torch
  9. from models.yolo import Model
  10. from utils.google_utils import attempt_download
  11. def create(name, pretrained, channels, classes):
  12. """Creates a specified YOLOv5 model
  13. Arguments:
  14. name (str): name of model, i.e. 'yolov5s'
  15. pretrained (bool): load pretrained weights into the model
  16. channels (int): number of input channels
  17. classes (int): number of model classes
  18. Returns:
  19. pytorch model
  20. """
  21. config = os.path.join(os.path.dirname(__file__), 'models', '%s.yaml' % name) # model.yaml path
  22. try:
  23. model = Model(config, channels, classes)
  24. if pretrained:
  25. ckpt = '%s.pt' % name # checkpoint filename
  26. attempt_download(ckpt) # download if not found locally
  27. state_dict = torch.load(ckpt, map_location=torch.device('cpu'))['model'].float().state_dict() # to FP32
  28. state_dict = {k: v for k, v in state_dict.items() if model.state_dict()[k].shape == v.shape} # filter
  29. model.load_state_dict(state_dict, strict=False) # load
  30. # model = model.autoshape() # cv2/PIL/np/torch inference: predictions = model(Image.open('image.jpg'))
  31. return model
  32. except Exception as e:
  33. help_url = 'https://github.com/ultralytics/yolov5/issues/36'
  34. s = 'Cache maybe be out of date, try force_reload=True. See %s for help.' % help_url
  35. raise Exception(s) from e
  36. def yolov5s(pretrained=False, channels=3, classes=80):
  37. """YOLOv5-small model from https://github.com/ultralytics/yolov5
  38. Arguments:
  39. pretrained (bool): load pretrained weights into the model, default=False
  40. channels (int): number of input channels, default=3
  41. classes (int): number of model classes, default=80
  42. Returns:
  43. pytorch model
  44. """
  45. return create('yolov5s', pretrained, channels, classes)
  46. def yolov5m(pretrained=False, channels=3, classes=80):
  47. """YOLOv5-medium model from https://github.com/ultralytics/yolov5
  48. Arguments:
  49. pretrained (bool): load pretrained weights into the model, default=False
  50. channels (int): number of input channels, default=3
  51. classes (int): number of model classes, default=80
  52. Returns:
  53. pytorch model
  54. """
  55. return create('yolov5m', pretrained, channels, classes)
  56. def yolov5l(pretrained=False, channels=3, classes=80):
  57. """YOLOv5-large model from https://github.com/ultralytics/yolov5
  58. Arguments:
  59. pretrained (bool): load pretrained weights into the model, default=False
  60. channels (int): number of input channels, default=3
  61. classes (int): number of model classes, default=80
  62. Returns:
  63. pytorch model
  64. """
  65. return create('yolov5l', pretrained, channels, classes)
  66. def yolov5x(pretrained=False, channels=3, classes=80):
  67. """YOLOv5-xlarge model from https://github.com/ultralytics/yolov5
  68. Arguments:
  69. pretrained (bool): load pretrained weights into the model, default=False
  70. channels (int): number of input channels, default=3
  71. classes (int): number of model classes, default=80
  72. Returns:
  73. pytorch model
  74. """
  75. return create('yolov5x', pretrained, channels, classes)