Du kannst nicht mehr als 25 Themen auswählen Themen müssen entweder mit einem Buchstaben oder einer Ziffer beginnen. Sie können Bindestriche („-“) enthalten und bis zu 35 Zeichen lang sein.

95 Zeilen
3.0KB

  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. model = Model(config, channels, classes)
  23. if pretrained:
  24. ckpt = '%s.pt' % name # checkpoint filename
  25. google_utils.attempt_download(ckpt) # download if not found locally
  26. state_dict = torch.load(ckpt, map_location=torch.device('cpu'))['model'].state_dict()
  27. state_dict = {k: v for k, v in state_dict.items() if model.state_dict()[k].numel() == v.numel()} # filter
  28. model.load_state_dict(state_dict, strict=False) # load
  29. return model
  30. def yolov5s(pretrained=False, channels=3, classes=80):
  31. """YOLOv5-small model from https://github.com/ultralytics/yolov5
  32. Arguments:
  33. pretrained (bool): load pretrained weights into the model, default=False
  34. channels (int): number of input channels, default=3
  35. classes (int): number of model classes, default=80
  36. Returns:
  37. pytorch model
  38. """
  39. return create('yolov5s', pretrained, channels, classes)
  40. def yolov5m(pretrained=False, channels=3, classes=80):
  41. """YOLOv5-medium model from https://github.com/ultralytics/yolov5
  42. Arguments:
  43. pretrained (bool): load pretrained weights into the model, default=False
  44. channels (int): number of input channels, default=3
  45. classes (int): number of model classes, default=80
  46. Returns:
  47. pytorch model
  48. """
  49. return create('yolov5m', pretrained, channels, classes)
  50. def yolov5l(pretrained=False, channels=3, classes=80):
  51. """YOLOv5-large model from https://github.com/ultralytics/yolov5
  52. Arguments:
  53. pretrained (bool): load pretrained weights into the model, default=False
  54. channels (int): number of input channels, default=3
  55. classes (int): number of model classes, default=80
  56. Returns:
  57. pytorch model
  58. """
  59. return create('yolov5l', pretrained, channels, classes)
  60. def yolov5x(pretrained=False, channels=3, classes=80):
  61. """YOLOv5-xlarge model from https://github.com/ultralytics/yolov5
  62. Arguments:
  63. pretrained (bool): load pretrained weights into the model, default=False
  64. channels (int): number of input channels, default=3
  65. classes (int): number of model classes, default=80
  66. Returns:
  67. pytorch model
  68. """
  69. return create('yolov5x', pretrained, channels, classes)