ultralytics 8.0.31 updates and fixes (#857)
Co-authored-by: Yonghye Kwon <developer.0hye@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Kalen Michael <kalenmike@gmail.com>
This commit is contained in:
Glenn Jocher
GitHub
parent
2e7a533ac3
commit
f5d003d05a
9 changed files with 285 additions and 131 deletions
|
|
@ -4,67 +4,62 @@ 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.utils import LOGGER, emojis
|
||||
from ultralytics.yolo.v8.detect import DetectionTrainer
|
||||
from ultralytics.hub.utils import split_key
|
||||
from ultralytics.yolo.engine.exporter import export_formats
|
||||
from ultralytics.yolo.engine.model import YOLO
|
||||
from ultralytics.yolo.utils import LOGGER, emojis, PREFIX
|
||||
|
||||
# Define all export formats
|
||||
EXPORT_FORMATS = list(export_formats()['Argument'][1:]) + ["ultralytics_tflite", "ultralytics_coreml"]
|
||||
|
||||
|
||||
def start(key=''):
|
||||
# Start training models with Ultralytics HUB. Usage: from src.ultralytics import start; start('API_KEY')
|
||||
def request_api_key(attempts=0):
|
||||
"""Prompt the user to input their API key"""
|
||||
import getpass
|
||||
|
||||
max_attempts = 3
|
||||
tries = f"Attempt {str(attempts + 1)} of {max_attempts}" if attempts > 0 else ""
|
||||
LOGGER.info(f"{PREFIX}Login. {tries}")
|
||||
input_key = getpass.getpass("Enter your Ultralytics HUB API key:\n")
|
||||
auth.api_key, model_id = split_key(input_key)
|
||||
if not auth.authenticate():
|
||||
attempts += 1
|
||||
LOGGER.warning(f"{PREFIX}Invalid API key ⚠️\n")
|
||||
if attempts < max_attempts:
|
||||
return request_api_key(attempts)
|
||||
raise ConnectionError(emojis(f"{PREFIX}Failed to authenticate ❌"))
|
||||
else:
|
||||
return model_id
|
||||
|
||||
def start(key=""):
|
||||
"""
|
||||
Start training models with Ultralytics HUB. Usage: from src.ultralytics import start; start('API_KEY')
|
||||
"""
|
||||
auth = Auth(key)
|
||||
try:
|
||||
api_key, model_id = split_key(key)
|
||||
auth = Auth(api_key) # attempts cookie login if no api key is present
|
||||
attempts = 1 if len(key) else 0
|
||||
if not auth.get_state():
|
||||
if len(key):
|
||||
LOGGER.warning(f"{PREFIX}Invalid API key ⚠️\n")
|
||||
model_id = request_api_key(attempts)
|
||||
LOGGER.info(f"{PREFIX}Authenticated ✅")
|
||||
model_id = request_api_key(auth)
|
||||
else:
|
||||
_, model_id = split_key(key)
|
||||
|
||||
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()
|
||||
|
||||
# TODO: refactor, hardcoded for v8
|
||||
args = session.model.copy()
|
||||
args.pop("id")
|
||||
args.pop("status")
|
||||
args.pop("weights")
|
||||
args["data"] = "coco128.yaml"
|
||||
args["model"] = "yolov8n.yaml"
|
||||
args["batch_size"] = 16
|
||||
args["imgsz"] = 64
|
||||
|
||||
trainer = DetectionTrainer(overrides=args)
|
||||
trainer = YOLO(session.input_file)
|
||||
session.register_callbacks(trainer)
|
||||
setattr(trainer, 'hub_session', session)
|
||||
trainer.train()
|
||||
trainer.train(**session.train_args)
|
||||
except Exception as e:
|
||||
LOGGER.warning(f"{PREFIX}{e}")
|
||||
|
||||
|
||||
def reset_model(key=''):
|
||||
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 HUB API key:\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})
|
||||
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")
|
||||
|
|
@ -72,38 +67,32 @@ def reset_model(key=''):
|
|||
LOGGER.warning(f"{PREFIX}model reset failure {r.status_code} {r.reason}")
|
||||
|
||||
|
||||
def export_model(key='', format='torchscript'):
|
||||
def export_model(key="", format="torchscript"):
|
||||
# Export a model to all formats
|
||||
assert format in EXPORT_FORMATS, f"Unsupported export format '{format}' passed, valid formats are {EXPORT_FORMATS}"
|
||||
api_key, model_id = split_key(key)
|
||||
formats = ('torchscript', 'onnx', 'openvino', 'engine', 'coreml', 'saved_model', 'pb', 'tflite', 'edgetpu', 'tfjs',
|
||||
'ultralytics_tflite', 'ultralytics_coreml')
|
||||
assert format in formats, f"ERROR: Unsupported export format '{format}' passed, valid formats are {formats}"
|
||||
|
||||
r = requests.post('https://api.ultralytics.com/export',
|
||||
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}"
|
||||
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'):
|
||||
def get_export(key="", format="torchscript"):
|
||||
# Get an exported model dictionary with download URL
|
||||
assert format in EXPORT_FORMATS, f"Unsupported export format '{format}' passed, valid formats are {EXPORT_FORMATS}"
|
||||
api_key, model_id = split_key(key)
|
||||
formats = ('torchscript', 'onnx', 'openvino', 'engine', 'coreml', 'saved_model', 'pb', 'tflite', 'edgetpu', 'tfjs',
|
||||
'ultralytics_tflite', 'ultralytics_coreml')
|
||||
assert format in formats, f"ERROR: Unsupported export format '{format}' passed, valid formats are {formats}"
|
||||
|
||||
r = requests.post('https://api.ultralytics.com/get-export',
|
||||
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}"
|
||||
assert (r.status_code == 200), f"{PREFIX}{format} get_export failure {r.status_code} {r.reason}"
|
||||
return r.json()
|
||||
|
||||
|
||||
# temp. For checking
|
||||
if __name__ == "__main__":
|
||||
start(key="b3fba421be84a20dbe68644e14436d1cce1b0a0aaa_HeMfHgvHsseMPhdq7Ylz")
|
||||
start()
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue