HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

Feature: Create HUB Models from CLI or Python Script (#7316)

Browse source 
Co-authored-by: Hassaan Farooq <103611273+hassaanfarooq01@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Kalen Michael 2024-01-10 02:36:14 +01:00 • committed by GitHub
parent a92adf8231
commit b54055a2c7
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 356 additions and 154 deletions
Download patch file Download diff file Expand all files Collapse all files

48
ultralytics/hub/__init__.py
View file

@ -1,28 +1,49 @@
# Ultralytics YOLO 🚀, AGPL-3.0 license
import requests
from hub_sdk import HUB_API_ROOT, HUB_WEB_ROOT, HUBClient
from ultralytics.data.utils import HUBDatasetStats
from ultralytics.hub.auth import Auth
from ultralytics.hub.utils import HUB_API_ROOT, HUB_WEB_ROOT, PREFIX
from ultralytics.hub.utils import PREFIX
from ultralytics.utils import LOGGER, SETTINGS
def login(api_key=''):
def login(api_key: str = None, save=True) -> bool:
"""
Log in to the Ultralytics HUB API using the provided API key.
The session is not stored; a new session is created when needed using the saved SETTINGS or the HUB_API_KEY environment variable if successfully authenticated.
Args:
api_key (str, optional): May be an API key or a combination API key and model ID, i.e. key_id
Example:
```python
from ultralytics import hub
hub.login('API_KEY')
```
api_key (str, optional): The API key to use for authentication. If not provided, it will be retrieved from SETTINGS or HUB_API_KEY environment variable.
save (bool, optional): Whether to save the API key to SETTINGS if authentication is successful.
Returns:
bool: True if authentication is successful, False otherwise.
"""
Auth(api_key, verbose=True)
api_key_url = f'{HUB_WEB_ROOT}/settings?tab=api+keys' # Set the redirect URL
saved_key = SETTINGS.get('api_key')
active_key = api_key or saved_key
credentials = {'api_key': active_key} if active_key and active_key != '' else None # Set credentials
client = HUBClient(credentials) # initialize HUBClient
if client.authenticated:
# Successfully authenticated with HUB
if save and client.api_key != saved_key:
SETTINGS.update({'api_key': client.api_key}) # update settings with valid API key
# Set message based on whether key was provided or retrieved from settings
log_message = ('New authentication successful ✅'
if client.api_key == api_key or not credentials else 'Authenticated ✅')
LOGGER.info(f'{PREFIX}{log_message}')
return True
else:
# Failed to authenticate with HUB
LOGGER.info(f'{PREFIX}Retrieve API key from {api_key_url}')
return False
def logout():
@ -43,7 +64,7 @@ def logout():
def reset_model(model_id=''):
"""Reset a trained model to an untrained state."""
r = requests.post(f'{HUB_API_ROOT}/model-reset', json={'apiKey': Auth().api_key, 'modelId': model_id})
r = requests.post(f'{HUB_API_ROOT}/model-reset', json={'modelId': model_id}, headers={'x-api-key': Auth().api_key})
if r.status_code == 200:
LOGGER.info(f'{PREFIX}Model reset successfully')
return
@ -73,7 +94,8 @@ def get_export(model_id='', format='torchscript'):
json={
'apiKey': Auth().api_key,
'modelId': model_id,
'format': format})
'format': format},
headers={'x-api-key': Auth().api_key})
assert r.status_code == 200, f'{PREFIX}{format} get_export failure {r.status_code} {r.reason}'
return r.json()