Nelze vybrat více než 25 témat Téma musí začínat písmenem nebo číslem, může obsahovat pomlčky („-“) a může být dlouhé až 35 znaků.

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798
  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 import google_utils
  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. google_utils.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. return model
  31. except Exception as e:
  32. help_url = 'https://github.com/ultralytics/yolov5/issues/36'
  33. print('%s\nCache maybe be out of date. Delete cache and retry. See %s for help.' % (e, help_url))
  34. def yolov5s(pretrained=False, channels=3, classes=80):
  35. """YOLOv5-small model from https://github.com/ultralytics/yolov5
  36. Arguments:
  37. pretrained (bool): load pretrained weights into the model, default=False
  38. channels (int): number of input channels, default=3
  39. classes (int): number of model classes, default=80
  40. Returns:
  41. pytorch model
  42. """
  43. return create('yolov5s', pretrained, channels, classes)
  44. def yolov5m(pretrained=False, channels=3, classes=80):
  45. """YOLOv5-medium model from https://github.com/ultralytics/yolov5
  46. Arguments:
  47. pretrained (bool): load pretrained weights into the model, default=False
  48. channels (int): number of input channels, default=3
  49. classes (int): number of model classes, default=80
  50. Returns:
  51. pytorch model
  52. """
  53. return create('yolov5m', pretrained, channels, classes)
  54. def yolov5l(pretrained=False, channels=3, classes=80):
  55. """YOLOv5-large model from https://github.com/ultralytics/yolov5
  56. Arguments:
  57. pretrained (bool): load pretrained weights into the model, default=False
  58. channels (int): number of input channels, default=3
  59. classes (int): number of model classes, default=80
  60. Returns:
  61. pytorch model
  62. """
  63. return create('yolov5l', pretrained, channels, classes)
  64. def yolov5x(pretrained=False, channels=3, classes=80):
  65. """YOLOv5-xlarge model from https://github.com/ultralytics/yolov5
  66. Arguments:
  67. pretrained (bool): load pretrained weights into the model, default=False
  68. channels (int): number of input channels, default=3
  69. classes (int): number of model classes, default=80
  70. Returns:
  71. pytorch model
  72. """
  73. return create('yolov5x', pretrained, channels, classes)