Spaces:
Build error
Build error
File size: 2,483 Bytes
2283b14 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
# Ultralytics YOLO 🚀, GPL-3.0 license
import requests
from ultralytics.hub.utils import HUB_API_ROOT, request_with_credentials
from ultralytics.yolo.utils import is_colab
API_KEY_PATH = "https://hub.ultralytics.com/settings?tab=api+keys"
class Auth:
id_token = api_key = model_key = False
def __init__(self, api_key=None):
self.api_key = self._clean_api_key(api_key)
self.authenticate() if self.api_key else self.auth_with_cookies()
@staticmethod
def _clean_api_key(key: str) -> str:
"""Strip model from key if present"""
separator = "_"
return key.split(separator)[0] if separator in key else key
def authenticate(self) -> bool:
"""Attempt to authenticate with server"""
try:
header = self.get_auth_header()
if header:
r = requests.post(f"{HUB_API_ROOT}/v1/auth", headers=header)
if not r.json().get('success', False):
raise ConnectionError("Unable to authenticate.")
return True
raise ConnectionError("User has not authenticated locally.")
except ConnectionError:
self.id_token = self.api_key = False # reset invalid
return False
def auth_with_cookies(self) -> bool:
"""
Attempt to fetch authentication via cookies and set id_token.
User must be logged in to HUB and running in a supported browser.
"""
if not is_colab():
return False # Currently only works with Colab
try:
authn = request_with_credentials(f"{HUB_API_ROOT}/v1/auth/auto")
if authn.get("success", False):
self.id_token = authn.get("data", {}).get("idToken", None)
self.authenticate()
return True
raise ConnectionError("Unable to fetch browser authentication details.")
except ConnectionError:
self.id_token = False # reset invalid
return False
def get_auth_header(self):
if self.id_token:
return {"authorization": f"Bearer {self.id_token}"}
elif self.api_key:
return {"x-api-key": self.api_key}
else:
return None
def get_state(self) -> bool:
"""Get the authentication state"""
return self.id_token or self.api_key
def set_api_key(self, key: str):
"""Get the authentication state"""
self.api_key = key
|