unitlab 2.0.4__tar.gz → 2.0.6__tar.gz

{unitlab-2.0.4/src/unitlab.egg-info → unitlab-2.0.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: unitlab
-Version: 2.0.4
+Version: 2.0.6
 Home-page: https://github.com/teamunitlab/unitlab-sdk
 Author: Unitlab Inc.
 Author-email: team@unitlab.ai

{unitlab-2.0.4 → unitlab-2.0.6}/setup.py

@@ -2,7 +2,7 @@ from setuptools import find_packages, setup
 
 setup(
     name="unitlab",
-    version="2.0.4",
+    version="2.0.6",
     license="MIT",
     author="Unitlab Inc.",
     author_email="team@unitlab.ai",

{unitlab-2.0.4 → unitlab-2.0.6}/src/unitlab/client.py

@@ -2,6 +2,7 @@ import asyncio
 import glob
 import logging
 import os
+import urllib.parse
 
 import aiofiles
 import aiohttp
@@ -9,12 +10,28 @@ import requests
 import tqdm
 
 from .dataset import DatasetUploadHandler
-from .exceptions import AuthenticationError
-from .utils import BASE_URL, ENDPOINTS, send_request
+from .exceptions import AuthenticationError, NetworkError
 
 logger = logging.getLogger(__name__)
 
 
+def handle_exceptions(f):
+    def throw_exception(*args, **kwargs):
+        try:
+            r = f(*args, **kwargs)
+            r_status = r.status_code
+            if r_status == 400:
+                raise AuthenticationError(
+                    "Please provide the api_key argument or set UNITLAB_API_KEY in your environment."
+                )
+            r.raise_for_status()
+            return r.json()
+        except requests.exceptions.RequestException as e:
+            raise NetworkError(str(e))
+
+    return throw_exception
+
+
 class UnitlabClient:
     """A client with a connection to the Unitlab.ai platform.
 
@@ -48,7 +65,7 @@ class UnitlabClient:
         :exc:`~unitlab.exceptions.AuthenticationError`: If an invalid API key is used or (when not passing the API key directly) if ``UNITLAB_API_KEY`` is not found in your environment.
     """
 
-    def __init__(self, api_key: str = None):
+    def __init__(self, api_key: str = None, api_url: str = None):
         if api_key is None:
             api_key = os.getenv("UNITLAB_API_KEY")
             if api_key is None:
@@ -56,7 +73,11 @@ class UnitlabClient:
                     message="Please provide the api_key argument or set UNITLAB_API_KEY in your environment."
                 )
             logger.info("Found a Unitlab API key in your environment.")
+        if api_url is None:
+            api_url = os.environ.get("UNITLAB_BASE_URL", "https://api.unitlab.ai")
+
         self.api_key = api_key
+        self.api_url = api_url
         self.api_session = requests.Session()
         adapter = requests.adapters.HTTPAdapter(max_retries=3)
         self.api_session.mount("http://", adapter)
@@ -96,152 +117,111 @@ class UnitlabClient:
     def _get_headers(self):
         return {"Authorization": f"Api-Key {self.api_key}"} if self.api_key else None
 
-    def projects(self):
-        response = send_request(
-            {
-                "method": "GET",
-                "endpoint": ENDPOINTS["projects"],
-                "headers": self._get_headers(),
-            },
-            session=self.api_session,
-        )
-        return response.json()
-
-    def project(self, project_id):
-        response = send_request(
-            {
-                "method": "GET",
-                "endpoint": ENDPOINTS["project"].format(project_id),
-                "headers": self._get_headers(),
-            },
-            session=self.api_session,
+    @handle_exceptions
+    def _get(self, endpoint):
+        return self.api_session.get(
+            urllib.parse.urljoin(self.api_url, endpoint), headers=self._get_headers()
         )
-        return response.json()
-
-    def project_members(self, project_id):
-        response = send_request(
-            {
-                "method": "GET",
-                "endpoint": ENDPOINTS["project_members"].format(project_id),
-                "headers": self._get_headers(),
-            },
-            session=self.api_session,
+
+    @handle_exceptions
+    def _post(self, endpoint, data):
+        return self.api_session.post(
+            urllib.parse.urljoin(self.api_url, endpoint),
+            json=data,
+            headers=self._get_headers(),
         )
-        return response.json()
+
+    def projects(self, pretty=0):
+        return self._get(f"/api/sdk/projects/?pretty={pretty}")
+
+    def project(self, project_id, pretty=0):
+        return self._get(f"/api/sdk/projects/{project_id}/?pretty={pretty}")
+
+    def project_members(self, project_id, pretty=0):
+        return self._get(f"/api/sdk/projects/{project_id}/members/?pretty={pretty}")
 
     def upload_data(self, project_id, directory, batch_size=100):
         if not os.path.isdir(directory):
             raise ValueError(f"Directory {directory} does not exist")
 
-        # set correct host
-        send_request(
-            {
-                "method": "GET",
-                "endpoint": ENDPOINTS["check"],
-                "headers": self._get_headers(),
-            },
-            session=self.api_session,
-        )
-        URL = os.environ["UNITLAB_BASE_URL"] + ENDPOINTS["upload_data"]
+        files = [
+            file
+            for files_list in (
+                glob.glob(os.path.join(directory, "") + extension)
+                for extension in ["*jpg", "*png", "*jpeg", "*webp"]
+            )
+            for file in files_list
+        ]
+        filtered_files = []
+        for file in files:
+            file_size = os.path.getsize(file) / 1024 / 1024
+            if file_size > 6:
+                logger.warning(
+                    f"File {file} is too large ({file_size:.4f} megabytes) skipping, max size is 6 MB"
+                )
+                continue
+            filtered_files.append(file)
+
+        num_files = len(filtered_files)
+        num_batches = (num_files + batch_size - 1) // batch_size
 
         async def post_file(
             session: aiohttp.ClientSession, file: str, project_id: str, retries=3
         ):
-            for _ in range(retries):
-                try:
-                    with open(file, "rb") as f:
-                        response = await session.request(
-                            "POST",
-                            url=URL,
-                            data=aiohttp.FormData(
-                                fields={"project": project_id, "file": f}
-                            ),
-                        )
-                        response.raise_for_status()
-                        return 1 if response.status == 201 else 0
-                except aiohttp.client_exceptions.ServerDisconnectedError as e:
-                    logger.warning(f"Error: {e}: Retrying...")
-                    await asyncio.sleep(0.1)
-                    continue
-                except Exception as e:
-                    logger.error(f"Error uploading file {file} - {e}")
-                    return 0
-
-        async def batch_upload(
-            session: aiohttp.ClientSession,
-            batch: list,
-            project_id: str,
-            pbar: tqdm.tqdm,
-        ):
-            tasks = []
-            for file in batch:
-                tasks.append(
-                    post_file(session=session, file=file, project_id=project_id)
+            async with aiofiles.open(file, "rb") as f:
+                form_data = aiohttp.FormData()
+                form_data.add_field("project", project_id)
+                form_data.add_field(
+                    "file", await f.read(), filename=os.path.basename(file)
                 )
-            for f in asyncio.as_completed(tasks):
-                pbar.update(await f)
+                for _ in range(retries):
+                    try:
+                        await asyncio.sleep(0.1)
+                        async with session.post(
+                            urllib.parse.urljoin(self.api_url, "/api/sdk/upload-data/"),
+                            data=form_data,
+                        ) as response:
+                            response.raise_for_status()
+                            return 1
+                    except aiohttp.ServerDisconnectedError as e:
+                        logger.warning(f"Server disconnected: {e}, Retrying...")
+                        await asyncio.sleep(0.1)
+                        continue
+                    except Exception as e:
+                        logger.error(f"Error uploading file {file} - {e}")
+                        return 0
 
         async def main():
-            files = [
-                file
-                for files_list in (
-                    glob.glob(os.path.join(directory, "") + extension)
-                    for extension in ["*jpg", "*png", "*jpeg", "*webp"]
-                )
-                for file in files_list
-            ]
-            filtered_files = []
-            for file in files:
-                file_size = os.path.getsize(file) / 1024 / 1024
-                if file_size > 6:
-                    logger.warning(
-                        f"File {file} is too large ({file_size:.4f} megabytes) skipping, max size is 6 MB"
-                    )
-                    continue
-                filtered_files.append(file)
-
-            num_files = len(filtered_files)
-            num_batches = (num_files + batch_size - 1) // batch_size
-
             logger.info(f"Uploading {num_files} files to project {project_id}")
             with tqdm.tqdm(total=num_files, ncols=80) as pbar:
                 async with aiohttp.ClientSession(
                     headers=self._get_headers()
                 ) as session:
                     for i in range(num_batches):
-                        await batch_upload(
-                            session,
-                            filtered_files[
-                                i * batch_size : min((i + 1) * batch_size, num_files)
-                            ],
-                            project_id,
-                            pbar,
-                        )
+                        tasks = []
+                        for file in filtered_files[
+                            i * batch_size : min((i + 1) * batch_size, num_files)
+                        ]:
+                            tasks.append(
+                                post_file(
+                                    session=session, file=file, project_id=project_id
+                                )
+                            )
+                        for f in asyncio.as_completed(tasks):
+                            pbar.update(await f)
 
         asyncio.run(main())
 
-    def datasets(self):
-        response = send_request(
-            {
-                "method": "GET",
-                "endpoint": ENDPOINTS["datasets"],
-                "headers": self._get_headers(),
-            },
-            session=self.api_session,
-        )
-        return response.json()
+    def datasets(self, pretty=0):
+        return self._get(f"/api/sdk/datasets/?pretty={pretty}")
 
     def dataset_download(self, dataset_id, export_type):
-        response = send_request(
-            {
-                "method": "POST",
-                "endpoint": ENDPOINTS["dataset"].format(dataset_id),
-                "headers": self._get_headers(),
-                "json": {"download_type": "annotation", "export_type": export_type},
-            },
-            session=self.api_session,
+        response = self._post(
+            f"/api/sdk/datasets/{dataset_id}/",
+            data={"download_type": "annotation", "export_type": export_type},
         )
-        with self.api_session.get(url=response.json()["file"], stream=True) as r:
+
+        with self.api_session.get(url=response["file"], stream=True) as r:
             r.raise_for_status()
             filename = f"dataset-{dataset_id}.json"
             with open(filename, "wb") as f:
@@ -251,20 +231,14 @@ class UnitlabClient:
             return os.path.abspath(filename)
 
     def download_dataset_files(self, dataset_id):
-        response = send_request(
-            {
-                "method": "POST",
-                "endpoint": ENDPOINTS["dataset"].format(dataset_id),
-                "headers": self._get_headers(),
-                "json": {"download_type": "files"},
-            },
-            session=self.api_session,
+        response = self._post(
+            f"/api/sdk/datasets/{dataset_id}/", data={"download_type": "files"}
         )
         folder = f"dataset-files-{dataset_id}"
         os.makedirs(folder, exist_ok=True)
         dataset_files = [
             dataset_file
-            for dataset_file in response.json()
+            for dataset_file in response
             if not os.path.isfile(os.path.join(folder, dataset_file["file_name"]))
         ]
 
@@ -285,22 +259,21 @@ class UnitlabClient:
                     return 1
 
         async def main():
-            async with aiohttp.ClientSession() as session:
-                tasks = [
-                    download_file(session=session, dataset_file=dataset_file)
-                    for dataset_file in dataset_files
-                ]
-                with tqdm.tqdm(total=len(dataset_files), ncols=80) as pbar:
+            with tqdm.tqdm(total=len(dataset_files), ncols=80) as pbar:
+                async with aiohttp.ClientSession() as session:
+                    tasks = [
+                        download_file(session=session, dataset_file=dataset_file)
+                        for dataset_file in dataset_files
+                    ]
                     for f in asyncio.as_completed(tasks):
                         pbar.update(await f)
 
         asyncio.run(main())
 
     def create_dataset(self, name, annotation_type, categories):
-        response = self.api_session.post(
-            url=f"{BASE_URL}/api/sdk/datasets/create/",
-            headers=self._get_headers(),
-            json={
+        response = self._post(
+            "/api/sdk/datasets/create/",
+            data={
                 "name": name,
                 "annotation_type": annotation_type,
                 "classes": [
@@ -309,35 +282,29 @@ class UnitlabClient:
                 ],
             },
         )
-        response.raise_for_status()
-        response = response.json()
         return response["pk"]
 
     def dataset_upload(
-        self, name, annotation_type, annotation_path, data_path, batch_size=100
+        self, name, annotation_type, annotation_path, data_path, batch_size=15
     ):
-        import random
-
         handler = DatasetUploadHandler(annotation_type, annotation_path, data_path)
         dataset_id = self.create_dataset(name, annotation_type, handler.categories)
-        img_ids = handler.getImgIds()
-        random.shuffle(img_ids)
-        image_ids = img_ids[:1000]
-        num_batches = (len(image_ids) + batch_size - 1) // batch_size
+        image_ids = handler.getImgIds()
+        url = urllib.parse.urljoin(
+            self.api_url, f"/api/sdk/datasets/{dataset_id}/upload/"
+        )
 
         async def main():
             with tqdm.tqdm(total=len(image_ids), ncols=80) as pbar:
                 async with aiohttp.ClientSession(
                     headers=self._get_headers()
                 ) as session:
-                    for i in range(num_batches):
+                    for i in range((len(image_ids) + batch_size - 1) // batch_size):
                         tasks = []
                         for image_id in image_ids[
                             i * batch_size : min((i + 1) * batch_size, len(image_ids))
                         ]:
-                            tasks.append(
-                                handler.upload_image(session, dataset_id, image_id)
-                            )
+                            tasks.append(handler.upload_image(session, url, image_id))
                         for f in asyncio.as_completed(tasks):
                             pbar.update(await f)
 

{unitlab-2.0.4 → unitlab-2.0.6}/src/unitlab/dataset.py

@@ -1,3 +1,4 @@
+import asyncio
 import itertools
 import json
 import logging
@@ -7,8 +8,6 @@ from collections import defaultdict
 import aiofiles
 import aiohttp
 
-from .utils import BASE_URL
-
 logger = logging.getLogger(__name__)
 
 
@@ -132,7 +131,7 @@ class COCO:
                     if ann["area"] > areaRng[0] and ann["area"] < areaRng[1]
                 ]
             )
-        if not iscrowd == None:
+        if iscrowd:
             ids = [ann["id"] for ann in anns if ann["iscrowd"] == iscrowd]
         else:
             ids = [ann["id"] for ann in anns]
@@ -201,7 +200,7 @@ class COCO:
         """
         if self._is_array_like(ids):
             return [self.anns[id] for id in ids]
-        elif type(ids) == int:
+        elif isinstance(ids, int):
             return [self.anns[ids]]
 
     def loadCats(self, ids=[]):
@@ -212,7 +211,7 @@ class COCO:
         """
         if self._is_array_like(ids):
             return [self.cats[id] for id in ids]
-        elif type(ids) == int:
+        elif isinstance(ids, int):
             return [self.cats[ids]]
 
     def loadImgs(self, ids=[]):
@@ -223,7 +222,7 @@ class COCO:
         """
         if self._is_array_like(ids):
             return [self.imgs[id] for id in ids]
-        elif type(ids) == int:
+        elif isinstance(ids, int):
             return [self.imgs[ids]]
 
 
@@ -278,26 +277,29 @@ class DatasetUploadHandler(COCO):
         if len(anns) == 0:
             logger.warning("No annotations found for image: {}".format(img_id))
             return
-        return self.get_img_bbox_payload(anns)
+        return getattr(self, f"get_{self.annotation_type}_payload")(anns)
 
-    async def upload_image(self, session, dataset_id, image_id):
+    async def upload_image(self, session, url, image_id, retries=3):
         image = self.loadImgs(image_id)[0]
         file_name = image["file_name"]
         payload = self.get_payload(image_id)
         if payload:
-            try:
-                async with aiofiles.open(
-                    os.path.join(self.data_path, file_name), "rb"
-                ) as f:
-                    form_data = aiohttp.FormData()
-                    form_data.add_field("file", await f.read(), filename=file_name)
-                    form_data.add_field("result", self.get_payload(image_id))
-                    async with session.post(
-                        f"{BASE_URL}/api/sdk/datasets/{dataset_id}/upload/",
-                        data=form_data,
-                    ) as response:
-                        response.raise_for_status()
-                        return 1
-            except Exception as e:
-                logger.error(f"Error uploading file {file_name} - {e}")
+            async with aiofiles.open(
+                os.path.join(self.data_path, file_name), "rb"
+            ) as f:
+                form_data = aiohttp.FormData()
+                form_data.add_field("file", await f.read(), filename=file_name)
+                form_data.add_field("result", self.get_payload(image_id))
+                for _ in range(retries):
+                    try:
+                        await asyncio.sleep(0.1)
+                        async with session.post(url, data=form_data) as response:
+                            response.raise_for_status()
+                            return 1
+                    except aiohttp.ServerDisconnectedError as e:
+                        logger.warning(f"Server disconnected - {e}, retrying...")
+                        await asyncio.sleep(0.1)
+                        continue
+                    except Exception as e:
+                        logger.error(f"Error uploading file {file_name} - {e}")
         return 0

{unitlab-2.0.4 → unitlab-2.0.6}/src/unitlab/exceptions.py

@@ -23,3 +23,9 @@ class AuthenticationError(UnitlabError):
     """Raised when an API key fails authentication."""
 
     pass
+
+
+class NetworkError(UnitlabError):
+    """Raised when an HTTP error occurs."""
+
+    pass

{unitlab-2.0.4 → unitlab-2.0.6}/src/unitlab/main.py

@@ -1,13 +1,11 @@
 from enum import Enum
 from pathlib import Path
-from typing import Optional
 from uuid import UUID
 
 import typer
 from typing_extensions import Annotated
 
 from .client import UnitlabClient
-from .utils import ENDPOINTS, send_request
 
 app = typer.Typer()
 project_app = typer.Typer()
@@ -25,10 +23,10 @@ class DownloadType(str, Enum):
 
 
 class AnnotationType(str, Enum):
-    img_bbox = "img_bbox"
-    img_polygon = "img_polygon"
-    img_semantic_segmentation = "img_semantic_segmentation"
-    img_skeleton = "img_skeleton"
+    IMG_BBOX = "img_bbox"
+    IMG_POLYGON = "img_polygon"
+    IMG_SEMANTIC_SEGMENTATION = "img_semantic_segmentation"
+    IMG_SKELETON = "img_skeleton"
 
 
 def get_client(api_key: str) -> UnitlabClient:
@@ -41,38 +39,17 @@ def get_headers(api_key):
 
 @project_app.command(name="list", help="Project list")
 def project_list(api_key: API_KEY):
-    response = send_request(
-        {
-            "method": "GET",
-            "headers": get_headers(api_key),
-            "endpoint": ENDPOINTS["cli_projects"],
-        }
-    )
-    print(response.json())
+    print(get_client(api_key).projects(pretty=1))
 
 
 @project_app.command(name="detail", help="Project detail")
 def project_detail(pk: UUID, api_key: API_KEY):
-    response = send_request(
-        {
-            "method": "GET",
-            "headers": get_headers(api_key),
-            "endpoint": ENDPOINTS["cli_project"].format(pk),
-        }
-    )
-    print(response.json())
+    print(get_client(api_key).project(project_id=pk, pretty=1))
 
 
 @project_app.command(help="Project members")
 def members(pk: UUID, api_key: API_KEY):
-    response = send_request(
-        {
-            "method": "GET",
-            "headers": get_headers(api_key),
-            "endpoint": ENDPOINTS["cli_project_members"].format(pk),
-        }
-    )
-    print(response.json())
+    print(get_client(api_key).project_members(project_id=pk, pretty=1))
 
 
 @project_app.command(help="Upload data")
@@ -94,14 +71,7 @@ if __name__ == "__main__":
 def dataset_list(
     api_key: API_KEY,
 ):
-    response = send_request(
-        {
-            "method": "GET",
-            "headers": get_headers(api_key),
-            "endpoint": ENDPOINTS["cli_datasets"],
-        }
-    )
-    print(response.json())
+    print(get_client(api_key).datasets(pretty=1))
 
 
 @dataset_app.command(name="upload", help="Upload dataset")
@@ -118,12 +88,9 @@ def dataset_upload(
     data_path: Annotated[
         Path, typer.Option(help="Directory containing the data to be uploaded")
     ],
-    batch_size: Annotated[
-        int, typer.Option(help="Batch size for uploading images")
-    ] = 100,
 ):
     get_client(api_key).dataset_upload(
-        name, annotation_type, annotation_path, data_path, batch_size
+        name, annotation_type.value, annotation_path, data_path
     )
 
 

{unitlab-2.0.4 → unitlab-2.0.6/src/unitlab.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: unitlab
-Version: 2.0.4
+Version: 2.0.6
 Home-page: https://github.com/teamunitlab/unitlab-sdk
 Author: Unitlab Inc.
 Author-email: team@unitlab.ai

{unitlab-2.0.4 → unitlab-2.0.6}/src/unitlab.egg-info/SOURCES.txt

@@ -8,7 +8,6 @@ src/unitlab/client.py
 src/unitlab/dataset.py
 src/unitlab/exceptions.py
 src/unitlab/main.py
-src/unitlab/utils.py
 src/unitlab.egg-info/PKG-INFO
 src/unitlab.egg-info/SOURCES.txt
 src/unitlab.egg-info/dependency_links.txt

unitlab-2.0.4/src/unitlab/utils.py

@@ -1,52 +0,0 @@
-import os
-
-import requests
-
-from .exceptions import AuthenticationError
-
-ENDPOINTS = {
-    "check": "/api/check/",
-    "projects": "/api/sdk/projects/",
-    "project": "/api/sdk/projects/{}/",
-    "project_members": "/api/sdk/projects/{}/members/",
-    "upload_data": "/api/sdk/upload-data/",
-    "datasets": "/api/sdk/datasets/",
-    "dataset": "/api/sdk/datasets/{}/",
-    "cli_projects": "/api/cli/projects/",
-    "cli_project": "/api/cli/projects/{}/",
-    "cli_project_members": "/api/cli/projects/{}/members/",
-    "cli_datasets": "/api/cli/datasets/",
-}
-BASE_URL = os.environ.get("UNITLAB_BASE_URL", "https://api.unitlab.ai")
-
-
-def send_request(request, session=None):
-    endpoint = request.pop("endpoint")
-    if os.environ.get("UNITLAB_BASE_URL"):
-        request["url"] = os.environ.get("UNITLAB_BASE_URL") + endpoint
-        response = (
-            session.request(**request) if session else requests.request(**request)
-        )
-        if response.ok:
-            return response
-        if response.status_code == 401:
-            raise AuthenticationError("Invalid API key")
-        elif response.status_code == 400:
-            raise Exception(response.json())
-        response.raise_for_status()
-        return response
-    else:
-        request["url"] = "https://api.unitlab.ai" + endpoint
-        response = (
-            session.request(**request) if session else requests.request(**request)
-        )
-        if response.ok:
-            os.environ["UNITLAB_BASE_URL"] = "https://api.unitlab.ai"
-            return response
-
-        if response.status_code == 401:
-            raise AuthenticationError("Invalid API key")
-        elif response.status_code == 400:
-            raise Exception(response.json())
-        response.raise_for_status()
-        return response