mirror of
https://github.com/THU-MIG/yolov10.git
synced 2025-05-23 13:34:23 +08:00
e7876e1ba9
Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: St. HeMeow <sheng.heyang@gmail.com> Co-authored-by: Danny Kim <imbird0312@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Co-authored-by: Torge Kummerow <CySlider@users.noreply.github.com> Co-authored-by: dankernel <dkdkernel@gmail.com> Co-authored-by: Burhan <62214284+Burhan-Q@users.noreply.github.com> Co-authored-by: Roshanlal <roshanlaladchitre103@gmail.com> Co-authored-by: Lorenzo Mammana <lorenzo.mammana@orobix.com> Co-authored-by: Yonghye Kwon <developer.0hye@gmail.com>
93 lines
3.4 KiB
Python
93 lines
3.4 KiB
Python
# Ultralytics YOLO 🚀, GPL-3.0 license
|
|
|
|
import requests
|
|
|
|
from ultralytics.hub.auth import Auth
|
|
from ultralytics.hub.session import HUBTrainingSession
|
|
from ultralytics.hub.utils import PREFIX, split_key
|
|
from ultralytics.yolo.engine.model import YOLO
|
|
from ultralytics.yolo.utils import LOGGER, emojis
|
|
|
|
|
|
def start(key=''):
|
|
"""
|
|
Start training models with Ultralytics HUB. Usage: from ultralytics.hub import start; start('API_KEY')
|
|
"""
|
|
auth = Auth(key)
|
|
model_id = split_key(key)[1] if auth.get_state() else request_api_key(auth)
|
|
if not model_id:
|
|
raise ConnectionError(emojis('Connecting with global API key is not currently supported. ❌'))
|
|
|
|
session = HUBTrainingSession(model_id=model_id, auth=auth)
|
|
session.check_disk_space()
|
|
|
|
model = YOLO(model=session.model_file, session=session)
|
|
model.train(**session.train_args)
|
|
|
|
|
|
def request_api_key(auth, max_attempts=3):
|
|
"""
|
|
Prompt the user to input their API key. Returns the model ID.
|
|
"""
|
|
import getpass
|
|
for attempts in range(max_attempts):
|
|
LOGGER.info(f'{PREFIX}Login. Attempt {attempts + 1} of {max_attempts}')
|
|
input_key = getpass.getpass(
|
|
'Enter your Ultralytics API Key from https://hub.ultralytics.com/settings?tab=api+keys:\n')
|
|
auth.api_key, model_id = split_key(input_key)
|
|
|
|
if auth.authenticate():
|
|
LOGGER.info(f'{PREFIX}Authenticated ✅')
|
|
return model_id
|
|
|
|
LOGGER.warning(f'{PREFIX}Invalid API key ⚠️\n')
|
|
|
|
raise ConnectionError(emojis(f'{PREFIX}Failed to authenticate ❌'))
|
|
|
|
|
|
def reset_model(key=''):
|
|
# Reset a trained model to an untrained state
|
|
api_key, model_id = split_key(key)
|
|
r = requests.post('https://api.ultralytics.com/model-reset', json={'apiKey': api_key, 'modelId': model_id})
|
|
|
|
if r.status_code == 200:
|
|
LOGGER.info(f'{PREFIX}Model reset successfully')
|
|
return
|
|
LOGGER.warning(f'{PREFIX}Model reset failure {r.status_code} {r.reason}')
|
|
|
|
|
|
def export_fmts_hub():
|
|
# Returns a list of HUB-supported export formats
|
|
from ultralytics.yolo.engine.exporter import export_formats
|
|
return list(export_formats()['Argument'][1:]) + ['ultralytics_tflite', 'ultralytics_coreml']
|
|
|
|
|
|
def export_model(key='', format='torchscript'):
|
|
# Export a model to all formats
|
|
assert format in export_fmts_hub(), f"Unsupported export format '{format}', valid formats are {export_fmts_hub()}"
|
|
api_key, model_id = split_key(key)
|
|
r = requests.post('https://api.ultralytics.com/export',
|
|
json={
|
|
'apiKey': api_key,
|
|
'modelId': model_id,
|
|
'format': format})
|
|
assert r.status_code == 200, f'{PREFIX}{format} export failure {r.status_code} {r.reason}'
|
|
LOGGER.info(f'{PREFIX}{format} export started ✅')
|
|
|
|
|
|
def get_export(key='', format='torchscript'):
|
|
# Get an exported model dictionary with download URL
|
|
assert format in export_fmts_hub, f"Unsupported export format '{format}', valid formats are {export_fmts_hub}"
|
|
api_key, model_id = split_key(key)
|
|
r = requests.post('https://api.ultralytics.com/get-export',
|
|
json={
|
|
'apiKey': api_key,
|
|
'modelId': model_id,
|
|
'format': format})
|
|
assert r.status_code == 200, f'{PREFIX}{format} get_export failure {r.status_code} {r.reason}'
|
|
return r.json()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
start()
|