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.

143 lines
5.7KB

  1. """YOLOv5 PyTorch Hub models https://pytorch.org/hub/ultralytics_yolov5/
  2. Usage:
  3. import torch
  4. model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
  5. """
  6. from pathlib import Path
  7. import torch
  8. from models.yolo import Model
  9. from utils.general import check_requirements, set_logging
  10. from utils.google_utils import attempt_download
  11. from utils.torch_utils import select_device
  12. dependencies = ['torch', 'yaml']
  13. check_requirements(Path(__file__).parent / 'requirements.txt', exclude=('pycocotools', 'thop'))
  14. def create(name, pretrained, channels, classes, autoshape, verbose):
  15. """Creates a specified YOLOv5 model
  16. Arguments:
  17. name (str): name of model, i.e. 'yolov5s'
  18. pretrained (bool): load pretrained weights into the model
  19. channels (int): number of input channels
  20. classes (int): number of model classes
  21. Returns:
  22. pytorch model
  23. """
  24. try:
  25. set_logging(verbose=verbose)
  26. cfg = list((Path(__file__).parent / 'models').rglob(f'{name}.yaml'))[0] # model.yaml path
  27. model = Model(cfg, channels, classes)
  28. if pretrained:
  29. fname = f'{name}.pt' # checkpoint filename
  30. attempt_download(fname) # download if not found locally
  31. ckpt = torch.load(fname, map_location=torch.device('cpu')) # load
  32. msd = model.state_dict() # model state_dict
  33. csd = ckpt['model'].float().state_dict() # checkpoint state_dict as FP32
  34. csd = {k: v for k, v in csd.items() if msd[k].shape == v.shape} # filter
  35. model.load_state_dict(csd, strict=False) # load
  36. if len(ckpt['model'].names) == classes:
  37. model.names = ckpt['model'].names # set class names attribute
  38. if autoshape:
  39. model = model.autoshape() # for file/URI/PIL/cv2/np inputs and NMS
  40. device = select_device('0' if torch.cuda.is_available() else 'cpu') # default to GPU if available
  41. return model.to(device)
  42. except Exception as e:
  43. help_url = 'https://github.com/ultralytics/yolov5/issues/36'
  44. s = 'Cache maybe be out of date, try force_reload=True. See %s for help.' % help_url
  45. raise Exception(s) from e
  46. def custom(path_or_model='path/to/model.pt', autoshape=True, verbose=True):
  47. """YOLOv5-custom model https://github.com/ultralytics/yolov5
  48. Arguments (3 options):
  49. path_or_model (str): 'path/to/model.pt'
  50. path_or_model (dict): torch.load('path/to/model.pt')
  51. path_or_model (nn.Module): torch.load('path/to/model.pt')['model']
  52. Returns:
  53. pytorch model
  54. """
  55. set_logging(verbose=verbose)
  56. model = torch.load(path_or_model) if isinstance(path_or_model, str) else path_or_model # load checkpoint
  57. if isinstance(model, dict):
  58. model = model['ema' if model.get('ema') else 'model'] # load model
  59. hub_model = Model(model.yaml).to(next(model.parameters()).device) # create
  60. hub_model.load_state_dict(model.float().state_dict()) # load state_dict
  61. hub_model.names = model.names # class names
  62. if autoshape:
  63. hub_model = hub_model.autoshape() # for file/URI/PIL/cv2/np inputs and NMS
  64. device = select_device('0' if torch.cuda.is_available() else 'cpu') # default to GPU if available
  65. return hub_model.to(device)
  66. def yolov5s(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
  67. # YOLOv5-small model https://github.com/ultralytics/yolov5
  68. return create('yolov5s', pretrained, channels, classes, autoshape, verbose)
  69. def yolov5m(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
  70. # YOLOv5-medium model https://github.com/ultralytics/yolov5
  71. return create('yolov5m', pretrained, channels, classes, autoshape, verbose)
  72. def yolov5l(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
  73. # YOLOv5-large model https://github.com/ultralytics/yolov5
  74. return create('yolov5l', pretrained, channels, classes, autoshape, verbose)
  75. def yolov5x(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
  76. # YOLOv5-xlarge model https://github.com/ultralytics/yolov5
  77. return create('yolov5x', pretrained, channels, classes, autoshape, verbose)
  78. def yolov5s6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
  79. # YOLOv5-small-P6 model https://github.com/ultralytics/yolov5
  80. return create('yolov5s6', pretrained, channels, classes, autoshape, verbose)
  81. def yolov5m6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
  82. # YOLOv5-medium-P6 model https://github.com/ultralytics/yolov5
  83. return create('yolov5m6', pretrained, channels, classes, autoshape, verbose)
  84. def yolov5l6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
  85. # YOLOv5-large-P6 model https://github.com/ultralytics/yolov5
  86. return create('yolov5l6', pretrained, channels, classes, autoshape, verbose)
  87. def yolov5x6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
  88. # YOLOv5-xlarge-P6 model https://github.com/ultralytics/yolov5
  89. return create('yolov5x6', pretrained, channels, classes, autoshape, verbose)
  90. if __name__ == '__main__':
  91. model = create(name='yolov5s', pretrained=True, channels=3, classes=80, autoshape=True, verbose=True) # pretrained
  92. # model = custom(path_or_model='path/to/model.pt') # custom
  93. # Verify inference
  94. import cv2
  95. import numpy as np
  96. from PIL import Image
  97. imgs = ['data/images/zidane.jpg', # filename
  98. 'https://github.com/ultralytics/yolov5/releases/download/v1.0/zidane.jpg', # URI
  99. cv2.imread('data/images/bus.jpg')[:, :, ::-1], # OpenCV
  100. Image.open('data/images/bus.jpg'), # PIL
  101. np.zeros((320, 640, 3))] # numpy
  102. results = model(imgs) # batched inference
  103. results.print()
  104. results.save()