__init__.py 3.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100
  1. # Ultralytics YOLO 🚀, AGPL-3.0 license
  2. import requests
  3. from ultralytics.data.utils import HUBDatasetStats
  4. from ultralytics.hub.auth import Auth
  5. from ultralytics.hub.utils import HUB_API_ROOT, HUB_WEB_ROOT, PREFIX
  6. from ultralytics.utils import LOGGER, SETTINGS
  7. def login(api_key=''):
  8. """
  9. Log in to the Ultralytics HUB API using the provided API key.
  10. Args:
  11. api_key (str, optional): May be an API key or a combination API key and model ID, i.e. key_id
  12. Example:
  13. ```python
  14. from ultralytics import hub
  15. hub.login('API_KEY')
  16. ```
  17. """
  18. Auth(api_key, verbose=True)
  19. def logout():
  20. """
  21. Log out of Ultralytics HUB by removing the API key from the settings file. To log in again, use 'yolo hub login'.
  22. Example:
  23. ```python
  24. from ultralytics import hub
  25. hub.logout()
  26. ```
  27. """
  28. SETTINGS['api_key'] = ''
  29. SETTINGS.save()
  30. LOGGER.info(f"{PREFIX}logged out ✅. To log in again, use 'yolo hub login'.")
  31. def reset_model(model_id=''):
  32. """Reset a trained model to an untrained state."""
  33. r = requests.post(f'{HUB_API_ROOT}/model-reset', json={'apiKey': Auth().api_key, 'modelId': model_id})
  34. if r.status_code == 200:
  35. LOGGER.info(f'{PREFIX}Model reset successfully')
  36. return
  37. LOGGER.warning(f'{PREFIX}Model reset failure {r.status_code} {r.reason}')
  38. def export_fmts_hub():
  39. """Returns a list of HUB-supported export formats."""
  40. from ultralytics.engine.exporter import export_formats
  41. return list(export_formats()['Argument'][1:]) + ['ultralytics_tflite', 'ultralytics_coreml']
  42. def export_model(model_id='', format='torchscript'):
  43. """Export a model to all formats."""
  44. assert format in export_fmts_hub(), f"Unsupported export format '{format}', valid formats are {export_fmts_hub()}"
  45. r = requests.post(f'{HUB_API_ROOT}/v1/models/{model_id}/export',
  46. json={'format': format},
  47. headers={'x-api-key': Auth().api_key})
  48. assert r.status_code == 200, f'{PREFIX}{format} export failure {r.status_code} {r.reason}'
  49. LOGGER.info(f'{PREFIX}{format} export started ✅')
  50. def get_export(model_id='', format='torchscript'):
  51. """Get an exported model dictionary with download URL."""
  52. assert format in export_fmts_hub(), f"Unsupported export format '{format}', valid formats are {export_fmts_hub()}"
  53. r = requests.post(f'{HUB_API_ROOT}/get-export',
  54. json={
  55. 'apiKey': Auth().api_key,
  56. 'modelId': model_id,
  57. 'format': format})
  58. assert r.status_code == 200, f'{PREFIX}{format} get_export failure {r.status_code} {r.reason}'
  59. return r.json()
  60. def check_dataset(path='', task='detect'):
  61. """
  62. Function for error-checking HUB dataset Zip file before upload. It checks a dataset for errors before it is
  63. uploaded to the HUB. Usage examples are given below.
  64. Args:
  65. path (str, optional): Path to data.zip (with data.yaml inside data.zip). Defaults to ''.
  66. task (str, optional): Dataset task. Options are 'detect', 'segment', 'pose', 'classify'. Defaults to 'detect'.
  67. Example:
  68. ```python
  69. from ultralytics.hub import check_dataset
  70. check_dataset('path/to/coco8.zip', task='detect') # detect dataset
  71. check_dataset('path/to/coco8-seg.zip', task='segment') # segment dataset
  72. check_dataset('path/to/coco8-pose.zip', task='pose') # pose dataset
  73. ```
  74. """
  75. HUBDatasetStats(path=path, task=task).get_json()
  76. LOGGER.info(f'Checks completed correctly ✅. Upload this dataset to {HUB_WEB_ROOT}/datasets/.')