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.

hubconf.py 3.9KB

4 yıl önce
4 yıl önce
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119
  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. from pathlib import Path
  7. import torch
  8. from models.yolo import Model
  9. from utils.general import set_logging
  10. from utils.google_utils import attempt_download
  11. dependencies = ['torch', 'yaml']
  12. set_logging()
  13. def create(name, pretrained, channels, classes):
  14. """Creates a specified YOLOv5 model
  15. Arguments:
  16. name (str): name of model, i.e. 'yolov5s'
  17. pretrained (bool): load pretrained weights into the model
  18. channels (int): number of input channels
  19. classes (int): number of model classes
  20. Returns:
  21. pytorch model
  22. """
  23. config = Path(__file__).parent / 'models' / f'{name}.yaml' # model.yaml path
  24. try:
  25. model = Model(config, channels, classes)
  26. if pretrained:
  27. fname = f'{name}.pt' # checkpoint filename
  28. attempt_download(fname) # download if not found locally
  29. ckpt = torch.load(fname, map_location=torch.device('cpu')) # load
  30. state_dict = ckpt['model'].float().state_dict() # to FP32
  31. state_dict = {k: v for k, v in state_dict.items() if model.state_dict()[k].shape == v.shape} # filter
  32. model.load_state_dict(state_dict, strict=False) # load
  33. if len(ckpt['model'].names) == classes:
  34. model.names = ckpt['model'].names # set class names attribute
  35. # model = model.autoshape() # for PIL/cv2/np inputs and NMS
  36. return model
  37. except Exception as e:
  38. help_url = 'https://github.com/ultralytics/yolov5/issues/36'
  39. s = 'Cache maybe be out of date, try force_reload=True. See %s for help.' % help_url
  40. raise Exception(s) from e
  41. def yolov5s(pretrained=False, channels=3, classes=80):
  42. """YOLOv5-small model from https://github.com/ultralytics/yolov5
  43. Arguments:
  44. pretrained (bool): load pretrained weights into the model, default=False
  45. channels (int): number of input channels, default=3
  46. classes (int): number of model classes, default=80
  47. Returns:
  48. pytorch model
  49. """
  50. return create('yolov5s', pretrained, channels, classes)
  51. def yolov5m(pretrained=False, channels=3, classes=80):
  52. """YOLOv5-medium model from https://github.com/ultralytics/yolov5
  53. Arguments:
  54. pretrained (bool): load pretrained weights into the model, default=False
  55. channels (int): number of input channels, default=3
  56. classes (int): number of model classes, default=80
  57. Returns:
  58. pytorch model
  59. """
  60. return create('yolov5m', pretrained, channels, classes)
  61. def yolov5l(pretrained=False, channels=3, classes=80):
  62. """YOLOv5-large model from https://github.com/ultralytics/yolov5
  63. Arguments:
  64. pretrained (bool): load pretrained weights into the model, default=False
  65. channels (int): number of input channels, default=3
  66. classes (int): number of model classes, default=80
  67. Returns:
  68. pytorch model
  69. """
  70. return create('yolov5l', pretrained, channels, classes)
  71. def yolov5x(pretrained=False, channels=3, classes=80):
  72. """YOLOv5-xlarge model from https://github.com/ultralytics/yolov5
  73. Arguments:
  74. pretrained (bool): load pretrained weights into the model, default=False
  75. channels (int): number of input channels, default=3
  76. classes (int): number of model classes, default=80
  77. Returns:
  78. pytorch model
  79. """
  80. return create('yolov5x', pretrained, channels, classes)
  81. if __name__ == '__main__':
  82. model = create(name='yolov5s', pretrained=True, channels=3, classes=80) # example
  83. model = model.fuse().autoshape() # for PIL/cv2/np inputs and NMS
  84. # Verify inference
  85. from PIL import Image
  86. imgs = [Image.open(x) for x in Path('data/images').glob('*.jpg')]
  87. results = model(imgs)
  88. results.show()
  89. results.print()