diff --git a/Makefile b/Makefile index bf0dcee..60aed3d 100644 --- a/Makefile +++ b/Makefile @@ -71,8 +71,6 @@ check-release: ## Check if everything is ready for release @echo "3. Commit: git commit -m '[LABIMP-XXXX] Prepare release vX.X.X'" @echo "4. Push and create PR to main (patch) or develop (minor)" -integration-test: - $(PYTHON) -m pytest -v labellerr_integration_tests.py pre-commit-install: pip install pre-commit diff --git a/driver.py b/driver.py index dbf9b61..ac37e32 100644 --- a/driver.py +++ b/driver.py @@ -5,7 +5,7 @@ from labellerr.client import LabellerrClient from labellerr.core.datasets import LabellerrDataset, create_dataset -from labellerr.core.connectors import LabellerrS3Connection +from labellerr.core.connectors import LabellerrS3Connection, LabellerrConnection from labellerr.core.files import LabellerrFile from labellerr.core.projects import ( LabellerrProject, @@ -25,26 +25,6 @@ ) from labellerr.core.autolabel import LabellerrAutoLabel -# Set logging level to DEBUG -logging.basicConfig(level=logging.DEBUG) - -load_dotenv() - -API_KEY = os.getenv("API_KEY") -API_SECRET = os.getenv("API_SECRET") -CLIENT_ID = os.getenv("CLIENT_ID") - -if not all([API_KEY, API_SECRET, CLIENT_ID]): - raise ValueError( - "API_KEY, API_SECRET, and CLIENT_ID must be set in environment variables" - ) - -# Initialize client -client = LabellerrClient( - api_key=API_KEY, - api_secret=API_SECRET, - client_id=CLIENT_ID, -) # if os.getenv("CREATE_DATASET", "").lower() == "true": # from labellerr import schemas @@ -160,25 +140,25 @@ # name="Amazon S3 Export Test", # description="Amazon S3 Export Test", # ))) -project = LabellerrProject(client=client, project_id="") - -export = project.create_export( - export_config=CreateExportParams( - export_name="Amazon S3 Export Test", - export_description="Amazon S3 Export Test", - export_format="json", - statuses=["review"], - connection_id=os.getenv("AWS_EXPORT_CONNECTION_ID"), - export_destination="s3", - export_folder_path="", # pattern - bucket_name/path/to/folder/ - the last slash is important - ) -) +# project = LabellerrProject(client=client, project_id="") -print(f"Export created: {export.report_id}") -print(f"Current status: {export._status}") +# export = project.create_export( +# export_config=CreateExportParams( +# export_name="Amazon S3 Export Test", +# export_description="Amazon S3 Export Test", +# export_format="json", +# statuses=["review"], +# connection_id=os.getenv("AWS_EXPORT_CONNECTION_ID"), +# export_destination="s3", +# export_folder_path="", # pattern - bucket_name/path/to/folder/ - the last slash is important +# ) +# ) + +# print(f"Export created: {export.report_id}") +# print(f"Current status: {export._status}") # Uncomment to poll until completion: -final_status = export.status() -print(f"Final status: {final_status}") +# final_status = export.status() +# print(f"Final status: {final_status}") # print(autolabel.train(training_request=TrainingRequest(model_id="yolov11", hyperparameters=Hyperparameters(epochs=10), slice_id='', min_samples_per_class=100, job_name="Yolo V11 Training"))) # dataset = LabellerrDataset(client=client, dataset_id="") @@ -201,9 +181,9 @@ # print(file.file_data) # project = LabellerrProject(client=client, project_id="") -# res = project.upload_preannotations( +# res = project.upload_preannotation( # annotation_format="coco_json", annotation_file="horses_coco.json" -# ) +# ).result() # print(res) # print(LabellerrProject.list_all_projects(client=client)) diff --git a/drivers/connectors.py b/drivers/connectors.py new file mode 100644 index 0000000..f5d51d5 --- /dev/null +++ b/drivers/connectors.py @@ -0,0 +1,73 @@ +import logging +import os +from dotenv import load_dotenv +from labellerr.client import LabellerrClient + +# from labellerr.core.connectors import ( +# LabellerrConnection, +# list_connections, +# delete_connection, +# create_connection, +# AWSConnectionParams, +# LabellerrGCSConnection, +# ) +# from labellerr.core.schemas import ( +# ConnectionType, +# ConnectorType, +# DatasetDataType, +# GCSConnectionTestParams, +# GCSConnectionParams, +# ) + +# Set logging level to DEBUG +logging.basicConfig(level=logging.DEBUG) + +load_dotenv() + +API_KEY = os.getenv("API_KEY") +API_SECRET = os.getenv("API_SECRET") +CLIENT_ID = os.getenv("CLIENT_ID") + +if not all([API_KEY, API_SECRET, CLIENT_ID]): + raise ValueError( + "API_KEY, API_SECRET, and CLIENT_ID must be set in environment variables" + ) + +# Initialize client +client = LabellerrClient( + api_key=API_KEY, + api_secret=API_SECRET, + client_id=CLIENT_ID, +) + +# response = LabellerrGCSConnection.test_connection(client=client, params=GCSConnectionTestParams( +# svc_account_json='labellerr-dev.json', +# path="gs://aws-labellerr-public-datasets/coco2017", +# connection_type=ConnectionType._IMPORT, +# data_type=DatasetDataType.video, +# )) +# print(response) + +# response = create_connection(client=client, connector_type=ConnectorType._GCS, params=GCSConnectionParams( +# svc_account_json='labellerr-dev.json', +# path="gs://aws-labellerr-public-datasets/coco2017", +# connection_type=ConnectionType._IMPORT, +# data_type=DatasetDataType.image, name="GCS Import Test", description="GCS Import Test")) +# print(response) +# connection = LabellerrConnection(client=client, connection_id='8c3dc4b4-e701-4d22-add3-28abc33e13ef') +# response = connection.test(path="gs://aws-labellerr-public-datasets/coco2017", connection_type=ConnectionType._IMPORT, data_type=DatasetDataType.image) +# print(response) +# connection = create_connection(client=client, connector_type=ConnectorType._S3, params=AWSConnectionParams( +# aws_access_key=os.getenv("AWS_KEY"), +# aws_secrets_key=os.getenv("AWS_SECRET"), +# s3_path="s3://amazon-s3-sync-test/labellerr-processed/videos/", # this path is not part of the connection but needed to test the connection on the desired path. +# # This can be dynamically changed while using the connection for creating datasets. +# connection_type=ConnectionType._IMPORT, +# name="Amazon S3 Import Test", +# description="Amazon S3 Import Test", +# )) +# print(f"Connection created: {connection.connection_id}") +# connection = LabellerrConnection(client=client, connection_id='2a30c044-57f7-42c7-8290-bab3bbac0ebc') +# print('connection type', connection.connection_type) + +# print (connection.test(s3_path="s3://amazon-s3-sync-test/labellerr-processed/videos/datasets/", connection_type=ConnectionType._IMPORT)) diff --git a/drivers/datasets.py b/drivers/datasets.py new file mode 100644 index 0000000..365401f --- /dev/null +++ b/drivers/datasets.py @@ -0,0 +1,61 @@ +# from labellerr.core.datasets import ( +# create_dataset_from_local, +# create_dataset_from_connection, +# ) +import logging +import os +from dotenv import load_dotenv +from labellerr.client import LabellerrClient +from labellerr.core.schemas import DatasetConfig +from labellerr.core.datasets import LabellerrDataset + +# from labellerr.core.connectors import ( +# LabellerrConnection, +# list_connections, +# delete_connection, +# create_connection, +# AWSConnectionParams, +# LabellerrGCSConnection, +# ) +# from labellerr.core.schemas import ( +# ConnectionType, +# ConnectorType, +# DatasetDataType, +# GCSConnectionTestParams, +# GCSConnectionParams, +# ) + +# Set logging level to DEBUG +logging.basicConfig(level=logging.DEBUG) + +load_dotenv() + +API_KEY = os.getenv("API_KEY") +API_SECRET = os.getenv("API_SECRET") +CLIENT_ID = os.getenv("CLIENT_ID") + +if not all([API_KEY, API_SECRET, CLIENT_ID]): + raise ValueError( + "API_KEY, API_SECRET, and CLIENT_ID must be set in environment variables" + ) + +# Initialize client +client = LabellerrClient( + api_key=API_KEY, + api_secret=API_SECRET, + client_id=CLIENT_ID, +) +dataset_config = DatasetConfig( + dataset_name="test_dataset_from_local SDK", + dataset_description="test dataset description", + data_type="image", +) +# dataset = create_dataset_from_local( +# client=client, +# dataset_config=dataset_config, +# folder_to_upload='images_single', +# ) +dataset = LabellerrDataset( + client=client, dataset_id="455e3d45-55f9-436d-98c2-07a514b7894e" +) +print(dataset.files_count) diff --git a/drivers/projects.py b/drivers/projects.py new file mode 100644 index 0000000..8e07c39 --- /dev/null +++ b/drivers/projects.py @@ -0,0 +1,65 @@ +# from labellerr.core.projects import LabellerrProject +# from labellerr.core.annotation_templates import LabellerrAnnotationTemplate +import logging +import os +from dotenv import load_dotenv +from labellerr.client import LabellerrClient + +# from labellerr.core.schemas.annotation_templates import CreateTemplateParams, AnnotationQuestion, QuestionType +# from labellerr.core.annotation_templates import create_template +# from labellerr.core.schemas.projects import CreateProjectParams, RotationConfig +# from labellerr.core.projects import create_project +# from labellerr.core.datasets import LabellerrDataset + + +# Set logging level to DEBUG +logging.basicConfig(level=logging.DEBUG) + +load_dotenv() + +API_KEY = os.getenv("API_KEY") +API_SECRET = os.getenv("API_SECRET") +CLIENT_ID = os.getenv("CLIENT_ID") + +if not all([API_KEY, API_SECRET, CLIENT_ID]): + raise ValueError( + "API_KEY, API_SECRET, and CLIENT_ID must be set in environment variables" + ) + +# Initialize client +client = LabellerrClient( + api_key=API_KEY, + api_secret=API_SECRET, + client_id=CLIENT_ID, +) + +# project = LabellerrProject(client=client, project_id="rafaela_youngest_pike_23125") +# res = project.upload_preannotation(annotation_format="coco_json", annotation_file="/Users/Ximi-Hoque/Downloads/export_to_annotate_05_15.json").result() +# print(res) + +# annotation_template = LabellerrAnnotationTemplate(client=client, annotation_template_id="00016829-9051-46b1-96c6-3ec6763c342a") +# print(annotation_template.annotation_template_data) + +# res = create_template(client, CreateTemplateParams(template_name="test_template_1", data_type="image", +# questions=[ +# AnnotationQuestion(question_number=1, +# question="test_question", +# question_id="test_question_id", +# question_type=QuestionType.bounding_box, +# required=True, +# color="#FF4500", +# )])) +# print(res) + +# project = create_project(client, +# CreateProjectParams( +# project_name="test_project_via_sdk", +# data_type="image", +# rotations=RotationConfig(annotation_rotation_count=1, review_rotation_count=1, client_review_rotation_count=1), +# use_ai=False), +# datasets=[LabellerrDataset(client=client, dataset_id="ca298293-7f5e-4bdd-801f-8863a5ba458b")], +# annotation_template=res +# ) +# print (project.project_data) +# project = LabellerrProject(client=client, project_id="dinnie_confidential_lynx_20766") +# print(project.project_data) diff --git a/labellerr/core/annotation_templates/__init__.py b/labellerr/core/annotation_templates/__init__.py new file mode 100644 index 0000000..f89441b --- /dev/null +++ b/labellerr/core/annotation_templates/__init__.py @@ -0,0 +1,64 @@ +from .base import LabellerrAnnotationTemplate +from ..schemas.annotation_templates import CreateTemplateParams, QuestionType, Option +from .. import constants +from ..client import LabellerrClient +import uuid + +__all__ = [ + "LabellerrAnnotationTemplate", +] + +object_types = [ + QuestionType.bounding_box, + QuestionType.polygon, + QuestionType.polyline, + QuestionType.dot, +] + + +def create_template( + client: LabellerrClient, params: CreateTemplateParams +) -> LabellerrAnnotationTemplate: + """Create an annotation template""" + + unique_id = str(uuid.uuid4()) + for question in params.questions: + if question.question_type in object_types: + if not question.color: + raise ValueError( + "Color is required for bounding box, polygon, polyline, and dot questions" + ) + question.options = [Option(option_name=question.color)] + else: + if question.question_type != QuestionType.input and not question.options: + raise ValueError( + "Options are required for radio, boolean, select, dropdown, stt, imc questions" + ) + + # Convert questions to the expected format + questions_data = [] + for question in params.questions: + question_dict = question.model_dump() + # Convert enum to string value + question_dict["option_type"] = question.question_type.value + # Remove question_type as it's now option_type + question_dict.pop("question_type", None) + questions_data.append(question_dict) + + payload = {"templateName": params.template_name, "questions": questions_data} + url = ( + f"{constants.BASE_URL}/annotations/create_template?client_id={client.client_id}&data_type={params.data_type.value}" + f"&uuid={unique_id}" + ) + + response = client.make_request( + "POST", + url, + extra_headers={"content-type": "application/json"}, + json=payload, + request_id=unique_id, + ) + return LabellerrAnnotationTemplate( + client=client, + annotation_template_id=response.get("response", None).get("template_id"), + ) diff --git a/labellerr/core/annotation_templates/base.py b/labellerr/core/annotation_templates/base.py new file mode 100644 index 0000000..21a24af --- /dev/null +++ b/labellerr/core/annotation_templates/base.py @@ -0,0 +1,50 @@ +from .. import constants +from ..client import LabellerrClient +from ..exceptions import InvalidAnnotationTemplateError +import uuid + + +class LabellerrAnnotationTemplate: + @staticmethod + def get_annotation_template(client: "LabellerrClient", annotation_template_id: str): + """Get annotation template from Labellerr API""" + unique_id = str(uuid.uuid4()) + url = ( + f"{constants.BASE_URL}/annotations/get_template?template_id={annotation_template_id}&client_id={client.client_id}" + f"&uuid={unique_id}" + ) + + response = client.make_request( + "GET", + url, + extra_headers={"content-type": "application/json"}, + request_id=unique_id, + ) + return response.get("response", None) + + """Base class for all Labellerr projects with factory behavior""" + + def __new__(cls, client: "LabellerrClient", annotation_template_id: str): + # Validate that the annotation template exists before creating the instance + annotation_template_data = cls.get_annotation_template( + client, annotation_template_id + ) + + if not annotation_template_data or ( + isinstance(annotation_template_data, dict) and not annotation_template_data + ): + raise InvalidAnnotationTemplateError( + f"Annotation template with ID '{annotation_template_id}' does not exist or could not be retrieved." + ) + + # Create the instance only if validation passes + instance = super().__new__(cls) + # Store the data on the instance to avoid calling API again in __init__ + instance.__annotation_template_data = annotation_template_data + return instance + + def __init__(self, client: "LabellerrClient", annotation_template_id: str): + self.client = client + self.annotation_template_id = annotation_template_id + # Use the data already fetched in __new__ + self.annotation_template_data = self.__annotation_template_data diff --git a/labellerr/core/connectors/__init__.py b/labellerr/core/connectors/__init__.py index cda7749..bf1f298 100644 --- a/labellerr/core/connectors/__init__.py +++ b/labellerr/core/connectors/__init__.py @@ -1,6 +1,9 @@ +import uuid from typing import TYPE_CHECKING -from ...schemas import AWSConnectionParams +from .. import constants +from ...schemas import AWSConnectionParams, ConnectionType, ConnectorType +from ...schemas import GCSConnectionParams from .connections import LabellerrConnection from .gcs_connection import GCSConnection as LabellerrGCSConnection from .s3_connection import S3Connection as LabellerrS3Connection @@ -13,60 +16,70 @@ def create_connection( client: "LabellerrClient", - connector_type: str, - client_id: str, - connector_config: dict, -): - """ - Sets up cloud connector (GCP/AWS) for dataset creation using factory pattern. + connector_type: ConnectorType, + params: AWSConnectionParams | GCSConnectionParams, +) -> LabellerrS3Connection | LabellerrGCSConnection: + if connector_type == ConnectorType._S3: + return LabellerrS3Connection.create_connection(client, params) + elif connector_type == ConnectorType._GCS: + return LabellerrGCSConnection.create_connection(client, params) + else: + raise ValueError(f"Unsupported connector type: {connector_type}") + +def list_connections( + client: "LabellerrClient", + connector: ConnectorType, + connection_type: ConnectionType = None, +) -> list[LabellerrGCSConnection | LabellerrS3Connection]: + """ + Lists connections for a client :param client: LabellerrClient instance - :param connector_type: Type of connector ('gcp' or 'aws') - :param client_id: Client ID - :param connector_config: Configuration dictionary for the connector - :return: Connection ID (str) for quick connection or full response (dict) for full connection + :param connection_type: Type of connection (import/export) + :param connector: Optional connector type filter (s3, gcs, etc.) + :return: List of connections """ - import logging - from ..exceptions import InvalidConnectionError + unique_id = str(uuid.uuid4()) + url = f"{constants.BASE_URL}/connectors/connections/list" + + params = { + "client_id": client.client_id, + "uuid": unique_id, + "connector": connector, + } + if connection_type: + params["connection_type"] = connection_type + extra_headers = {"email_id": client.api_key} - try: - if connector_type == "gcs": - from .gcs_connection import GCSConnection + response = client.make_request( + "GET", url, extra_headers=extra_headers, request_id=unique_id, params=params + ) + return [ + LabellerrConnection(client, connection["connection_id"]) + for connection in response.get("response", []) + ] + + +def delete_connection(client: "LabellerrClient", connection_id: str): + """ + Deletes a connector connection by ID. + :param connection_id: The ID of the connection to delete + :return: Parsed JSON response + """ + request_id = str(uuid.uuid4()) + url = ( + f"{constants.BASE_URL}/connectors/connections/delete" + f"?client_id={client.client_id}&uuid={request_id}" + ) - return GCSConnection.create_connection(client, connector_config) - elif connector_type == "s3": - from .s3_connection import S3Connection + extra_headers = {"email_id": client.api_key} - # Determine which method to call based on config parameters - # Full connection has: aws_access_key, aws_secrets_key, s3_path, name, description - # Quick connection has: bucket_name, folder_path, access_key_id, secret_access_key - if "aws_access_key" in connector_config and "name" in connector_config: - # Full connection flow - creates a saved connection - return S3Connection.setup_full_connection( - client, - AWSConnectionParams( - client_id=connector_config["client_id"], - aws_access_key=connector_config["aws_access_key"], - aws_secrets_key=connector_config["aws_secrets_key"], - s3_path=connector_config["s3_path"], - data_type=connector_config["data_type"], - name=connector_config["name"], - description=connector_config["description"], - connection_type=connector_config.get( - "connection_type", "import" - ), - ), - ) - else: - # Quick connection flow - for dataset creation - return S3Connection.create_connection( - client, client_id, connector_config - ) - else: - raise InvalidConnectionError( - f"Unsupported connector type: {connector_type}" - ) - except Exception as e: - logging.error(f"Failed to setup {connector_type} connector: {e}") - raise + response = client.make_request( + "POST", + url, + extra_headers=extra_headers, + request_id=request_id, + json={"connection_id": connection_id}, + ) + return response.get("response", None) diff --git a/labellerr/core/connectors/connections.py b/labellerr/core/connectors/connections.py index b18a6d4..0c56f27 100644 --- a/labellerr/core/connectors/connections.py +++ b/labellerr/core/connectors/connections.py @@ -1,14 +1,13 @@ """This module will contain all CRUD for connections. Example, create, list connections, get connection, delete connection, update connection, etc.""" import uuid -from abc import ABCMeta, abstractmethod -from typing import TYPE_CHECKING, Dict - +from abc import ABCMeta +from typing import Dict from .. import client_utils, constants -from ..exceptions import InvalidConnectionError, InvalidDatasetIDError +from ..schemas import ConnectionType, DatasetDataType +from ..exceptions import InvalidConnectionError -if TYPE_CHECKING: - from ..client import LabellerrClient +from ..client import LabellerrClient class LabellerrConnectionMeta(ABCMeta): @@ -23,24 +22,19 @@ def _register(cls, connection_type, connection_class): @staticmethod def get_connection(client: "LabellerrClient", connection_id: str): """Get connection from Labellerr API""" - # ------------------------------- [needs refactoring after we consolidate api_calls into one function ] --------------------------------- + + assert connection_id, "Connection ID is can't be empty" unique_id = str(uuid.uuid4()) - url = ( - f"{constants.BASE_URL}/connections/{connection_id}?client_id={client.client_id}" - f"&uuid={unique_id}" - ) - headers = client_utils.build_headers( - api_key=client.api_key, - api_secret=client.api_secret, - client_id=client.client_id, - extra_headers={"content-type": "application/json"}, - ) + url = f"{constants.BASE_URL}/connectors/connections/{connection_id}/details" - response = client_utils.request( - "GET", url, headers=headers, request_id=unique_id + params = {"client_id": client.client_id, "uuid": unique_id} + + extra_headers = {"content-type": "application/json"} + + response = client.make_request( + "GET", url, extra_headers=extra_headers, request_id=unique_id, params=params ) return response.get("response", None) - # ------------------------------- [needs refactoring after we consolidate api_calls into one function ] --------------------------------- """Metaclass that combines ABC functionality with factory pattern""" @@ -54,16 +48,11 @@ def __call__(cls, client, connection_id, **kwargs): return instance connection_data = cls.get_connection(client, connection_id) if connection_data is None: - raise InvalidDatasetIDError(f"Connection not found: {connection_id}") - connection_type = connection_data.get("connection_type") - if connection_type not in constants.CONNECTION_TYPES: - raise InvalidConnectionError( - f"Connection type not supported: {connection_type}" - ) - - connection_class = cls._registry.get(connection_type) + raise InvalidConnectionError(f"Connection not found: {connection_id}") + connector = connection_data.get("connector") + connection_class = cls._registry.get(connector) if connection_class is None: - raise InvalidConnectionError(f"Unknown connection type: {connection_type}") + raise InvalidConnectionError(f"Unknown connector type: {connector}") kwargs["connection_data"] = connection_data return connection_class(client, connection_id, **kwargs) @@ -74,84 +63,65 @@ class LabellerrConnection(metaclass=LabellerrConnectionMeta): def __init__(self, client: "LabellerrClient", connection_id: str, **kwargs): self.client = client self._connection_id_input = connection_id - self.connection_data = kwargs["connection_data"] + self.__connection_data = kwargs["connection_data"] @property - def connection_id(self): - return self.connection_data.get("connection_id") + def name(self): + return self.__connection_data.get("name") @property - def connection_type(self): - return self.connection_data.get("connection_type") - - @abstractmethod - def test_connection(self): - """Each connection type must implement its own connection testing logic""" - pass - - def list_connections( - self, - connection_type: str, - connector: str = None, - ) -> list: - """ - List connections for a client - :param connection_type: Type of connection (import/export) - :param connector: Optional connector type filter (s3, gcs, etc.) - :return: List of connections - """ - request_uuid = str(uuid.uuid4()) - list_connection_url = ( - f"{constants.BASE_URL}/connectors/connections/list" - f"?client_id={self.client.client_id}&uuid={request_uuid}&connection_type={connection_type}" - ) - - if connector: - list_connection_url += f"&connector={connector}" + def description(self): + return self.__connection_data.get("description") - headers = client_utils.build_headers( - api_key=self.client.api_key, - api_secret=self.client.api_secret, - client_id=self.client.client_id, - extra_headers={"email_id": self.client.api_key}, - ) + @property + def connection_id(self): + return self.__connection_data.get("connection_id") - return client_utils.request( - "GET", list_connection_url, headers=headers, request_id=request_uuid - ) + @property + def connection_type(self): + return self.__connection_data.get("connection_type") - def delete_connection(self, connection_id: str): - """ - Deletes a connector connection by ID. - :param connection_id: The ID of the connection to delete - :return: Parsed JSON response - """ - import json + @property + def connector(self): + return self.__connection_data.get("connector") - from ... import schemas + @property + def created_at(self): + return self.__connection_data.get("created_at") - # Validate parameters using Pydantic - params = schemas.DeleteConnectionParams( - client_id=self.client.client_id, connection_id=connection_id - ) - request_uuid = str(uuid.uuid4()) - delete_url = ( - f"{constants.BASE_URL}/connectors/connections/delete" - f"?client_id={params.client_id}&uuid={request_uuid}" + @property + def created_by(self): + return self.__connection_data.get("created_by") + + def test( + self, path: str, connection_type: ConnectionType, data_type: DatasetDataType + ): + request_id = str(uuid.uuid4()) + test_connection_url = ( + f"{constants.BASE_URL}/connectors/connections/test" + f"?client_id={self.client.client_id}&uuid={request_id}" ) headers = client_utils.build_headers( api_key=self.client.api_key, api_secret=self.client.api_secret, client_id=self.client.client_id, - extra_headers={ - "content-type": "application/json", - "email_id": self.client.api_key, - }, + extra_headers={"email_id": self.client.api_key}, ) - payload = json.dumps({"connection_id": params.connection_id}) - - return client_utils.request( - "POST", delete_url, headers=headers, data=payload, request_id=request_uuid + # Test endpoint also expects multipart/form-data format + test_request = { + "connector": (None, self.connector), + "path": (None, path), + "connection_type": (None, connection_type.value), + "connection_id": (None, self.connection_id), + "data_type": (None, data_type.value), + } + response = client_utils.request( + "POST", + test_connection_url, + headers=headers, + files=test_request, + request_id=request_id, ) + return response.get("response", {}) diff --git a/labellerr/core/connectors/gcs_connection.py b/labellerr/core/connectors/gcs_connection.py index 6ff811f..da114b2 100644 --- a/labellerr/core/connectors/gcs_connection.py +++ b/labellerr/core/connectors/gcs_connection.py @@ -1,21 +1,72 @@ import uuid -from typing import TYPE_CHECKING - +import logging +import os from .. import client_utils, constants +from ..client import LabellerrClient from .connections import LabellerrConnection, LabellerrConnectionMeta - -if TYPE_CHECKING: - from labellerr import LabellerrClient +from ..schemas import GCSConnectionTestParams, GCSConnectionParams +from ..exceptions import LabellerrError class GCSConnection(LabellerrConnection): - def test_connection(self): - print("Testing GCS connection!") - return True + @staticmethod + def test_connection( + client: "LabellerrClient", params: GCSConnectionTestParams + ) -> dict: + request_id = str(uuid.uuid4()) + url = ( + f"{constants.BASE_URL}/connectors/connections/test" + f"?client_id={client.client_id}&uuid={request_id}" + ) + + # Prepare multipart form data + files = [] + if params.svc_account_json and os.path.exists(params.svc_account_json): + # If file path is provided, read and upload as file + with open(params.svc_account_json, "rb") as f: + files = [ + ( + "attachment_files", + ( + params.svc_account_json.split("/")[-1], + f.read(), + "application/json", + ), + ) + ] + else: + raise LabellerrError("Service account JSON file is required") + + # Prepare form data payload + form_data = { + "connector": "gcs", + "path": params.path, + "connection_type": params.connection_type.value, + "data_type": params.data_type.value, + } + + # Build headers without content-type for multipart form data + headers = client_utils.build_headers( + api_key=client.api_key, + api_secret=client.api_secret, + client_id=client.client_id, + ) + + response = client_utils.request( + "POST", + url, + headers=headers, + data=form_data, + files=files, + request_id=request_id, + ) + return response.get("response", {}) @staticmethod - def create_connection(client: "LabellerrClient", gcp_config: dict) -> str: + def create_connection( + client: "LabellerrClient", params: GCSConnectionParams + ) -> str: """ Sets up GCP connector for dataset creation (quick connection). @@ -23,37 +74,54 @@ def create_connection(client: "LabellerrClient", gcp_config: dict) -> str: :param gcp_config: GCP configuration containing bucket_name, folder_path, service_account_key :return: Connection ID for GCP connector """ - import json - - from ... import LabellerrError - - required_fields = ["bucket_name"] - for field in required_fields: - if field not in gcp_config: - raise LabellerrError(f"Required field '{field}' missing in gcp_config") + response = GCSConnection.test_connection(client, params) + logging.info(f"GCS connection test response: {response}") unique_id = str(uuid.uuid4()) - url = f"{constants.BASE_URL}/connectors/connect/gcp?client_id={client.client_id}&uuid={unique_id}" + url = f"{constants.BASE_URL}/connectors/connections/create?client_id={client.client_id}&uuid={unique_id}" + # Build headers without content-type for multipart form data headers = client_utils.build_headers( api_key=client.api_key, api_secret=client.api_secret, client_id=client.client_id, - extra_headers={"content-type": "application/json"}, ) - payload = json.dumps( - { - "bucket_name": gcp_config["bucket_name"], - "folder_path": gcp_config.get("folder_path", ""), - "service_account_key": gcp_config.get("service_account_key"), - } - ) + # Prepare multipart form data + with open(params.svc_account_json, "rb") as f: + files = [ + ( + "attachment_files", + ( + params.svc_account_json.split("/")[-1], + f.read(), + "application/json", + ), + ) + ] + + # Prepare form data (not JSON) + form_data = { + "connector": "gcs", + "connection_type": params.connection_type.value, + "name": params.name, + "description": params.description, + "credentials": "svc_account_json", + "client_id": client.client_id, + } response_data = client_utils.request( - "POST", url, headers=headers, data=payload, request_id=unique_id + "POST", + url, + headers=headers, + data=form_data, + files=files, + request_id=unique_id, + ) + return LabellerrConnection( + client=client, + connection_id=response_data.get("response", {}).get("connection_id"), ) - return response_data["response"]["connection_id"] LabellerrConnectionMeta._register("gcs", GCSConnection) diff --git a/labellerr/core/connectors/s3_connection.py b/labellerr/core/connectors/s3_connection.py index 0e8662c..8aa15b4 100644 --- a/labellerr/core/connectors/s3_connection.py +++ b/labellerr/core/connectors/s3_connection.py @@ -1,13 +1,10 @@ import json import uuid -from typing import TYPE_CHECKING - +from ..client import LabellerrClient from ..schemas import AWSConnectionParams, AWSConnectionTestParams from .. import client_utils, constants from .connections import LabellerrConnection, LabellerrConnectionMeta - -if TYPE_CHECKING: - from labellerr import LabellerrClient +import logging class S3Connection(LabellerrConnection): @@ -15,17 +12,17 @@ class S3Connection(LabellerrConnection): @staticmethod def test_connection( client: "LabellerrClient", params: AWSConnectionTestParams - ) -> bool: + ) -> dict: """ Tests an AWS S3 connection. :param client: The LabellerrClient instance :param params: The AWS connection parameters :return: True if the connection is successful, False otherwise """ - request_uuid = str(uuid.uuid4()) + request_id = str(uuid.uuid4()) test_connection_url = ( f"{constants.BASE_URL}/connectors/connections/test" - f"?client_id={client.client_id}&uuid={request_uuid}" + f"?client_id={client.client_id}&uuid={request_id}" ) headers = client_utils.build_headers( @@ -46,8 +43,9 @@ def test_connection( test_request = { "credentials": (None, aws_credentials_json), "connector": (None, "s3"), - "path": (None, params.s3_path), + "path": (None, params.path), "connection_type": (None, params.connection_type.value), + "data_type": (None, params.data_type.value), } # Remove content-type from headers to let requests set it with boundary @@ -55,14 +53,14 @@ def test_connection( k: v for k, v in headers.items() if k.lower() != "content-type" } - client_utils.request( + response = client_utils.request( "POST", test_connection_url, headers=headers_without_content_type, files=test_request, - request_id=request_uuid, + request_id=request_id, ) - return True + return response.get("response", {}) @staticmethod def create_connection( @@ -76,6 +74,11 @@ def create_connection( :return: Dictionary containing the response from the API """ + # Tests the connection before creating it + response = S3Connection.test_connection(client, params) + + logging.info(f"Connection test response: {response}") + unique_id = str(uuid.uuid4()) url = f"{constants.BASE_URL}/connectors/connections/create?client_id={client.client_id}&uuid={unique_id}" @@ -95,7 +98,6 @@ def create_connection( request_payload = { "credentials": (None, creds_payload), "connector": (None, "s3"), - "path": (None, params.s3_path), "connection_type": (None, params.connection_type.value), "name": (None, params.name), "description": (None, params.description), @@ -104,7 +106,10 @@ def create_connection( response_data = client_utils.request( "POST", url, headers=headers, files=request_payload, request_id=unique_id ) - return response_data.get("response", {}) + return LabellerrConnection( + client=client, + connection_id=response_data.get("response", {}).get("connection_id"), + ) LabellerrConnectionMeta._register("s3", S3Connection) diff --git a/labellerr/core/constants.py b/labellerr/core/constants.py index 91e1bf9..e35c554 100644 --- a/labellerr/core/constants.py +++ b/labellerr/core/constants.py @@ -1,4 +1,4 @@ -BASE_URL = "https://api.labellerr.com" +BASE_URL = "https://api-gateway-722091373895.us-central1.run.app" ALLOWED_ORIGINS = "https://pro.labellerr.com" @@ -7,7 +7,7 @@ TOTAL_FILES_SIZE_LIMIT_PER_DATASET = 2.5 * 1024 * 1024 * 1024 # 2.5GB TOTAL_FILES_COUNT_LIMIT_PER_DATASET = 2500 -ANNOTATION_FORMAT = ["json", "coco_json", "csv", "png"] +ANNOTATION_FORMAT = ["json", "coco_json", "csv", "png", "video_json", "video_json"] LOCAL_EXPORT_FORMAT = ["json", "coco_json", "csv", "png"] LOCAL_EXPORT_STATUS = [ "review", diff --git a/labellerr/core/datasets/__init__.py b/labellerr/core/datasets/__init__.py index 28c782c..4d975b7 100644 --- a/labellerr/core/datasets/__init__.py +++ b/labellerr/core/datasets/__init__.py @@ -1,11 +1,9 @@ import json import logging import uuid -from typing import TYPE_CHECKING -from ... import schemas as root_schemas from .. import constants, schemas -from ..connectors import create_connection +from ..client import LabellerrClient from ..exceptions import LabellerrError from .audio_dataset import AudioDataSet as LabellerrAudioDataset from .base import LabellerrDataset @@ -14,9 +12,6 @@ from .utils import upload_files, upload_folder_files_to_dataset from .video_dataset import VideoDataset as LabellerrVideoDataset -if TYPE_CHECKING: - from ..client import LabellerrClient - __all__ = [ "LabellerrImageDataset", "LabellerrVideoDataset", @@ -26,14 +21,51 @@ ] -def create_dataset( +def create_dataset_from_connection( + client: "LabellerrClient", + dataset_config: schemas.DatasetConfig, + connection_id: str, + path: str, +) -> LabellerrDataset: + """ + Creates a dataset via a connection. + + :param client: The client to use for the request. + :param dataset_config: The configuration for the dataset. + :param connection_id: The ID of the connection to use for the dataset. + :param path: The path to the data source. + :return: The LabellerrDataset instance. + """ + unique_id = str(uuid.uuid4()) + url = f"{constants.BASE_URL}/datasets/create?client_id={client.client_id}&uuid={unique_id}" + + payload = json.dumps( + { + "dataset_name": dataset_config.dataset_name, + "dataset_description": dataset_config.dataset_description, + "data_type": dataset_config.data_type, + "connection_id": connection_id, + "path": path, + "client_id": client.client_id, + "es_multimodal_index": dataset_config.multimodal_indexing, + } + ) + response_data = client.make_request( + "POST", + url, + extra_headers={"content-type": "application/json"}, + request_id=unique_id, + data=payload, + ) + dataset_id = response_data["response"]["dataset_id"] + return LabellerrDataset(client=client, dataset_id=dataset_id) # type: ignore[abstract] + + +def create_dataset_from_local( client: "LabellerrClient", dataset_config: schemas.DatasetConfig, files_to_upload=None, folder_to_upload=None, - path=None, - connection_id=None, - connector_config=None, ): """ Creates a dataset with support for multiple data types and connectors. @@ -44,129 +76,138 @@ def create_dataset( Can also be a DatasetConfig Pydantic model instance. :param files_to_upload: List of file paths to upload (for local connector) :param folder_to_upload: Path to folder to upload (for local connector) - :param connection_id: Pre-existing connection ID to use for the dataset. - Either connection_id or connector_config can be provided, but not both. - :param connector_config: Configuration for cloud connectors (GCP/AWS) - Can be a dict or AWSConnectorConfig/GCPConnectorConfig model instance. - Either connection_id or connector_config can be provided, but not both. - :return: A dictionary containing the response status and the ID of the created dataset. - :raises LabellerrError: If both connection_id and connector_config are provided. + :return: The LabellerrDataset instance. + """ + if files_to_upload is not None: + connection_id = upload_files( + client, + client_id=client.client_id, + files_list=files_to_upload, + ) + elif folder_to_upload is not None: + result = upload_folder_files_to_dataset( + client, + { + "client_id": client.client_id, + "folder_path": folder_to_upload, + "data_type": dataset_config.data_type, + }, + ) + connection_id = result.pop("connection_id") + logging.info(f"Folder uploaded successfully. {result}") + else: + raise LabellerrError("No files or folder to upload provided") + + return create_dataset_from_connection( + client, + dataset_config, + connection_id, + path="local", + ) + + +def delete_dataset(client: "LabellerrClient", dataset_id: str): + """ + Deletes a dataset from the system. + + :param dataset_id: The ID of the dataset to delete + :return: Dictionary containing deletion status + :raises LabellerrError: If the deletion fails """ + unique_id = str(uuid.uuid4()) + url = f"{constants.BASE_URL}/datasets/{dataset_id}/delete?client_id={client.client_id}&uuid={unique_id}" + + return client.make_request( + "DELETE", + url, + extra_headers={"content-type": "application/json"}, + request_id=unique_id, + ) - try: - # Validate that both connection_id and connector_config are not provided - if connection_id is not None and connector_config is not None: - raise LabellerrError( - "Cannot provide both connection_id and connector_config. " - "Use connection_id for existing connections or connector_config to create a new connection." + +def list_datasets( + client: "LabellerrClient", + datatype: str, + scope: schemas.DataSetScope, + page_size: int = constants.DEFAULT_PAGE_SIZE, + last_dataset_id: str = None, +): + """ + Retrieves datasets by parameters with pagination support. + Always returns a generator that yields individual datasets. + + :param client: The client object. + :param datatype: The type of data for the dataset. + :param scope: The permission scope for the dataset. + :param page_size: Number of datasets to return per page (default: 10) + Use -1 to auto-paginate through all pages + Use specific number to fetch only that many datasets from first page + :param last_dataset_id: ID of the last dataset from previous page for pagination + (only used when page_size is a specific number, ignored for -1) + :return: Generator yielding individual datasets + + Examples: + # Auto-paginate through all datasets + for dataset in get_all_datasets(client, "image", DataSetScope.client, page_size=-1): + print(dataset) + + # Get first 20 datasets + datasets = list(get_all_datasets(client, "image", DataSetScope.client, page_size=20)) + + # Manual pagination - first page of 10 + gen = get_all_datasets(client, "image", DataSetScope.client, page_size=10) + first_10 = list(gen) + """ + # Auto-pagination mode: yield datasets across all pages + if page_size == -1: + actual_page_size = constants.DEFAULT_PAGE_SIZE + current_last_dataset_id = None + has_more = True + + while has_more: + unique_id = str(uuid.uuid4()) + url = ( + f"{constants.BASE_URL}/datasets/list?client_id={client.client_id}&data_type={datatype}&permission_level={scope}" + f"&page_size={actual_page_size}&uuid={unique_id}" ) - connector_type = dataset_config.connector_type - # Use provided connection_id or set to None (will be created later if needed) - final_connection_id = connection_id - - # Handle different connector types only if connection_id is not provided - if final_connection_id is None: - if connector_type == "local": - if files_to_upload is not None: - try: - final_connection_id = upload_files( - client, - client_id=client.client_id, - files_list=files_to_upload, - ) - except Exception as e: - raise LabellerrError( - f"Failed to upload files to dataset: {str(e)}" - ) - - elif folder_to_upload is not None: - try: - result = upload_folder_files_to_dataset( - client, - { - "client_id": client.client_id, - "folder_path": folder_to_upload, - "data_type": dataset_config.data_type, - }, - ) - final_connection_id = result["connection_id"] - except Exception as e: - raise LabellerrError( - f"Failed to upload folder files to dataset: {str(e)}" - ) - elif connector_config is None: - # Create empty dataset for local connector - final_connection_id = None - - elif connector_type in ["gcp", "aws"]: - if connector_config is None: - raise LabellerrError( - f"connector_config is required for {connector_type} connector when connection_id is not provided" - ) - if path is None: - raise LabellerrError( - f"path is required for {connector_type} connector" - ) - - # Validate connector_config using Pydantic models - if connector_type == "aws": - if not isinstance( - connector_config, root_schemas.AWSConnectorConfig - ): - validated_connector = root_schemas.AWSConnectorConfig( - **connector_config - ) - else: - validated_connector = connector_config - else: # gcp - if not isinstance( - connector_config, root_schemas.GCPConnectorConfig - ): - validated_connector = root_schemas.GCPConnectorConfig( - **connector_config - ) - else: - validated_connector = connector_config - - try: - final_connection_id = create_connection( - client, - connector_type, - client.client_id, - validated_connector.model_dump(), - ) - except Exception as e: - raise LabellerrError( - f"Failed to setup {connector_type} connector: {str(e)}" - ) - else: - raise LabellerrError(f"Unsupported connector type: {connector_type}") + if current_last_dataset_id: + url += f"&last_dataset_id={current_last_dataset_id}" - unique_id = str(uuid.uuid4()) - url = f"{constants.BASE_URL}/datasets/create?client_id={client.client_id}&uuid={unique_id}" + response = client.make_request( + "GET", + url, + extra_headers={"content-type": "application/json"}, + request_id=unique_id, + ) - payload = json.dumps( - { - "dataset_name": dataset_config.dataset_name, - "dataset_description": dataset_config.dataset_description, - "data_type": dataset_config.data_type, - "connection_id": final_connection_id, - "path": path, - "client_id": client.client_id, - "connector_type": connector_type, - } + datasets = response.get("response", {}).get("datasets", []) + for dataset in datasets: + yield dataset + + # Check if there are more pages + has_more = response.get("response", {}).get("has_more", False) + current_last_dataset_id = response.get("response", {}).get( + "last_dataset_id" + ) + + else: + unique_id = str(uuid.uuid4()) + url = ( + f"{constants.BASE_URL}/datasets/list?client_id={client.client_id}&data_type={datatype}&permission_level={scope}" + f"&page_size={page_size}&uuid={unique_id}" ) - response_data = client.make_request( - "POST", + + # Add last_dataset_id for pagination if provided + if last_dataset_id: + url += f"&last_dataset_id={last_dataset_id}" + + response = client.make_request( + "GET", url, extra_headers={"content-type": "application/json"}, request_id=unique_id, - data=payload, ) - dataset_id = response_data["response"]["dataset_id"] - return LabellerrDataset(client=client, dataset_id=dataset_id) # type: ignore[abstract] - - except LabellerrError as e: - logging.error(f"Failed to create dataset: {e}") - raise + datasets = response.get("response", {}).get("datasets", []) + for dataset in datasets: + yield dataset diff --git a/labellerr/core/datasets/base.py b/labellerr/core/datasets/base.py index 177adc3..9360f32 100644 --- a/labellerr/core/datasets/base.py +++ b/labellerr/core/datasets/base.py @@ -1,16 +1,14 @@ """This module will contain all CRUD for datasets. Example, create, list datasets, get dataset, delete dataset, update dataset, etc.""" import json +import logging import uuid from abc import ABCMeta, abstractmethod -from typing import TYPE_CHECKING, Dict +from typing import Dict, Optional, Any -from ...schemas import DataSetScope from .. import constants from ..exceptions import InvalidDatasetError - -if TYPE_CHECKING: - from ..client import LabellerrClient +from ..client import LabellerrClient class LabellerrDatasetMeta(ABCMeta): @@ -53,8 +51,6 @@ def __call__(cls, client, dataset_id, **kwargs): if dataset_data is None: raise InvalidDatasetError(f"Dataset not found: {dataset_id}") data_type = dataset_data.get("data_type") - if data_type not in constants.DATA_TYPES: - raise InvalidDatasetError(f"Data type not supported: {data_type}") dataset_class = cls._registry.get(data_type) if dataset_class is None: @@ -69,136 +65,117 @@ class LabellerrDataset(metaclass=LabellerrDatasetMeta): def __init__(self, client: "LabellerrClient", dataset_id: str, **kwargs): self.client = client self.dataset_id = dataset_id - self.dataset_data = kwargs["dataset_data"] + self.__dataset_data = kwargs["dataset_data"] @property - def files_count(self): - return self.dataset_data.get("files_count", 0) + def name(self): + return self.__dataset_data.get("name") @property - def status_code(self): - return self.dataset_data.get("status_code", 501) # if not found, return 501 + def description(self): + return self.__dataset_data.get("description") @property - def data_type(self): - return self.dataset_data.get("data_type") + def created_at(self): + return self.__dataset_data.get("created_at") - @abstractmethod - def fetch_files(self): - """Each file type must implement its own download logic""" - pass + @property + def created_by(self): + return self.__dataset_data.get("created_by") - @staticmethod - def get_all_datasets( - client: "LabellerrClient", - datatype: str, - scope: DataSetScope, - page_size: int = None, - last_dataset_id: str = None, - ): - """ - Retrieves datasets by parameters with pagination support. - Always returns a generator that yields individual datasets. - - :param client: The client object. - :param datatype: The type of data for the dataset. - :param scope: The permission scope for the dataset. - :param page_size: Number of datasets to return per page (default: 10) - Use -1 to auto-paginate through all pages - Use specific number to fetch only that many datasets from first page - :param last_dataset_id: ID of the last dataset from previous page for pagination - (only used when page_size is a specific number, ignored for -1) - :return: Generator yielding individual datasets + @property + def files_count(self): + return self.__dataset_data.get("files_count", 0) - Examples: - # Auto-paginate through all datasets - for dataset in get_all_datasets(client, "image", DataSetScope.client, page_size=-1): - print(dataset) + @property + def status_code(self): + return self.__dataset_data.get("status_code", 501) # if not found, return 501 - # Get first 20 datasets - datasets = list(get_all_datasets(client, "image", DataSetScope.client, page_size=20)) + @property + def data_type(self): + return self.__dataset_data.get("data_type") - # Manual pagination - first page of 10 - gen = get_all_datasets(client, "image", DataSetScope.client, page_size=10) - first_10 = list(gen) + def status( + self, + interval: float = 2.0, + timeout: Optional[float] = None, + max_retries: Optional[int] = None, + ) -> Dict[str, Any]: """ - # Set default page size if not specified - if page_size is None: - page_size = constants.DEFAULT_PAGE_SIZE - - # Auto-pagination mode: yield datasets across all pages - if page_size == -1: - actual_page_size = constants.DEFAULT_PAGE_SIZE - current_last_dataset_id = None - has_more = True - - while has_more: - unique_id = str(uuid.uuid4()) - url = ( - f"{constants.BASE_URL}/datasets/list?client_id={client.client_id}&data_type={datatype}&permission_level={scope}" - f"&page_size={actual_page_size}&uuid={unique_id}" - ) + Poll dataset status until completion or timeout. - if current_last_dataset_id: - url += f"&last_dataset_id={current_last_dataset_id}" + Args: + interval: Time in seconds between status checks (default: 2.0) + timeout: Maximum time in seconds to poll before giving up + max_retries: Maximum number of retries before giving up - response = client.make_request( - "GET", - url, - extra_headers={"content-type": "application/json"}, - request_id=unique_id, - ) + Returns: + Final dataset data with status information - datasets = response.get("response", {}).get("datasets", []) - for dataset in datasets: - yield dataset + Examples: + # Poll until dataset processing is complete + final_status = dataset.status() - # Check if there are more pages - has_more = response.get("response", {}).get("has_more", False) - current_last_dataset_id = response.get("response", {}).get( - "last_dataset_id" - ) + # Poll with custom timeout + final_status = dataset.status(timeout=300) + + # Poll with custom interval and max retries + final_status = dataset.status(interval=5.0, max_retries=20) + """ + from ..utils import poll - else: + def get_dataset_status(): unique_id = str(uuid.uuid4()) url = ( - f"{constants.BASE_URL}/datasets/list?client_id={client.client_id}&data_type={datatype}&permission_level={scope}" - f"&page_size={page_size}&uuid={unique_id}" + f"{constants.BASE_URL}/datasets/{self.dataset_id}?client_id={self.client.client_id}" + f"&uuid={unique_id}" ) - # Add last_dataset_id for pagination if provided - if last_dataset_id: - url += f"&last_dataset_id={last_dataset_id}" - - response = client.make_request( + response = self.client.make_request( "GET", url, extra_headers={"content-type": "application/json"}, request_id=unique_id, ) - datasets = response.get("response", {}).get("datasets", []) - for dataset in datasets: - yield dataset - - def delete_dataset(self, dataset_id): - """ - Deletes a dataset from the system. - - :param dataset_id: The ID of the dataset to delete - :return: Dictionary containing deletion status - :raises LabellerrError: If the deletion fails - """ - unique_id = str(uuid.uuid4()) - url = f"{constants.BASE_URL}/datasets/{dataset_id}/delete?client_id={self.client.client_id}&uuid={unique_id}" - - return self.client.make_request( - "DELETE", - url, - extra_headers={"content-type": "application/json"}, - request_id=unique_id, + dataset_data = response.get("response", {}) + if dataset_data: + self.__dataset_data = dataset_data + return dataset_data + + def is_completed(dataset_data): + status_code = dataset_data.get("status_code", 500) + # Consider dataset complete when status_code is 200 (success) or >= 400 (error/failed) + return status_code == 200 or status_code >= 400 + + def on_success(dataset_data): + status_code = dataset_data.get("status_code", 500) + if status_code == 300: + logging.info( + "Dataset %s processing completed successfully!", self.dataset_id + ) + else: + logging.warning( + "Dataset %s processing finished with status code: %s", + self.dataset_id, + status_code, + ) + return dataset_data + + return poll( + function=get_dataset_status, + condition=is_completed, + interval=interval, + timeout=timeout, + max_retries=max_retries, + on_success=on_success, ) - def sync_datasets( + @abstractmethod + def fetch_files(self): + """Each file type must implement its own download logic""" + pass + + def sync_with_connection( self, project_id, path, diff --git a/labellerr/core/datasets/datasets.py b/labellerr/core/datasets/datasets.py deleted file mode 100644 index 886d9bf..0000000 --- a/labellerr/core/datasets/datasets.py +++ /dev/null @@ -1,764 +0,0 @@ -import json -import logging -import os -import uuid -from concurrent.futures import ThreadPoolExecutor, as_completed - -import requests - -from labellerr import client_utils, gcs, schemas, utils -from labellerr.core import constants -from labellerr.exceptions import LabellerrError -from labellerr.utils import validate_params - - -class DataSets(object): - """ - Handles dataset-related operations for the Labellerr API. - """ - - def __init__(self, api_key, api_secret, client): - """ - Initialize the DataSets handler. - - :param api_key: The API key for authentication - :param api_secret: The API secret for authentication - :param client: Reference to the parent Labellerr Client instance for delegating certain operations - """ - self.api_key = api_key - self.api_secret = api_secret - self.client = client - - def create_project( - self, - project_name, - data_type, - client_id, - attached_datasets, - annotation_template_id, - rotations, - use_ai=False, - created_by=None, - ): - """ - Creates a project with the given configuration. - - :param project_name: Name of the project - :param data_type: Type of data (image, video, etc.) - :param client_id: ID of the client - :param attached_datasets: List of dataset IDs to attach to the project - :param annotation_template_id: ID of the annotation template - :param rotations: Dictionary containing rotation configuration - :param use_ai: Boolean flag for AI usage (default: False) - :param created_by: Optional creator information - :return: Project creation response - :raises LabellerrError: If the creation fails - """ - # Validate parameters using Pydantic - params = schemas.CreateProjectParams( - project_name=project_name, - data_type=data_type, - client_id=client_id, - attached_datasets=attached_datasets, - annotation_template_id=annotation_template_id, - rotations=rotations, - use_ai=use_ai, - created_by=created_by, - ) - unique_id = str(uuid.uuid4()) - url = f"{constants.BASE_URL}/projects/create?client_id={params.client_id}&uuid={unique_id}" - - payload = json.dumps( - { - "project_name": params.project_name, - "attached_datasets": params.attached_datasets, - "data_type": params.data_type, - "annotation_template_id": str(params.annotation_template_id), - "rotations": params.rotations.model_dump(), - "use_ai": params.use_ai, - "created_by": params.created_by, - } - ) - - headers = client_utils.build_headers( - api_key=self.api_key, - api_secret=self.api_secret, - client_id=params.client_id, - extra_headers={ - "Origin": constants.ALLOWED_ORIGINS, - "Content-Type": "application/json", - }, - ) - - return client_utils.request( - "POST", url, headers=headers, data=payload, request_id=unique_id - ) - - def initiate_create_project(self, payload): - """ - Orchestrates project creation by handling dataset creation, annotation guidelines, - and final project setup. - """ - - try: - # validate all the parameters - required_params = [ - "client_id", - "dataset_name", - "dataset_description", - "data_type", - "created_by", - "project_name", - # Either annotation_guide or annotation_template_id must be provided - "autolabel", - ] - for param in required_params: - if param not in payload: - raise LabellerrError(f"Required parameter {param} is missing") - - if param == "client_id": - if ( - not isinstance(payload[param], str) - or not payload[param].strip() - ): - raise LabellerrError("client_id must be a non-empty string") - - # Validate created_by email format - created_by = payload.get("created_by") - if ( - not isinstance(created_by, str) - or "@" not in created_by - or "." not in created_by.split("@")[-1] - ): - raise LabellerrError("Please enter email id in created_by") - - # Ensure either annotation_guide or annotation_template_id is provided - if not payload.get("annotation_guide") and not payload.get( - "annotation_template_id" - ): - raise LabellerrError( - "Please provide either annotation guide or annotation template id" - ) - - # If annotation_guide is provided, validate its entries - if payload.get("annotation_guide"): - for guide in payload["annotation_guide"]: - if "option_type" not in guide: - raise LabellerrError( - "option_type is required in annotation_guide" - ) - if guide["option_type"] not in constants.OPTION_TYPE_LIST: - raise LabellerrError( - f"option_type must be one of {constants.OPTION_TYPE_LIST}" - ) - - if "folder_to_upload" in payload and "files_to_upload" in payload: - raise LabellerrError( - "Cannot provide both files_to_upload and folder_to_upload" - ) - - if "folder_to_upload" not in payload and "files_to_upload" not in payload: - raise LabellerrError( - "Either files_to_upload or folder_to_upload must be provided" - ) - - if ( - isinstance(payload.get("files_to_upload"), list) - and len(payload["files_to_upload"]) == 0 - ): - payload.pop("files_to_upload") - - if "rotation_config" not in payload: - payload["rotation_config"] = { - "annotation_rotation_count": 1, - "review_rotation_count": 1, - "client_review_rotation_count": 1, - } - self.validate_rotation_config(payload["rotation_config"]) - - if payload["data_type"] not in constants.DATA_TYPES: - raise LabellerrError( - f"Invalid data_type. Must be one of {constants.DATA_TYPES}" - ) - - logging.info("Rotation configuration validated . . .") - - logging.info("Creating dataset . . .") - dataset_response = self.create_dataset( - { - "client_id": payload["client_id"], - "dataset_name": payload["dataset_name"], - "data_type": payload["data_type"], - "dataset_description": payload["dataset_description"], - }, - files_to_upload=payload.get("files_to_upload"), - folder_to_upload=payload.get("folder_to_upload"), - ) - - dataset_id = dataset_response["dataset_id"] - - def check_dataset_status(): - dataset_status = self.client.get_dataset( - payload["client_id"], dataset_id - ) - - if isinstance(dataset_status, dict): - if "response" in dataset_status: - return dataset_status["response"].get("status_code", 200) == 300 - else: - return True - return False - - def is_dataset_ready(status): - return status is True - - def on_success(status): - logging.info("Dataset created and ready for use") - - def on_exception(e): - logging.error(f"Error checking dataset status: {str(e)}") - raise LabellerrError(f"Failed to check dataset status: {str(e)}") - - utils.poll( - function=check_dataset_status, - condition=is_dataset_ready, - interval=5, - on_success=on_success, - on_exception=on_exception, - ) - - if payload.get("annotation_template_id"): - annotation_template_id = payload["annotation_template_id"] - else: - annotation_template_id = self.create_annotation_guideline( - payload["client_id"], - payload["annotation_guide"], - payload["project_name"], - payload["data_type"], - ) - logging.info(f"Annotation guidelines created {annotation_template_id}") - - project_response = self.create_project( - project_name=payload["project_name"], - data_type=payload["data_type"], - client_id=payload["client_id"], - attached_datasets=[dataset_id], - annotation_template_id=annotation_template_id, - rotations=payload["rotation_config"], - use_ai=payload.get("use_ai", False), - created_by=payload["created_by"], - ) - - return { - "status": "success", - "message": "Project created successfully", - "project_id": project_response, - } - - except LabellerrError: - raise - except Exception: - logging.exception("Unexpected error in project creation") - raise - - def create_annotation_guideline( - self, client_id, questions, template_name, data_type - ): - """ - Updates the annotation guideline for a project. - - :param config: A dictionary containing the project ID, data type, client ID, autolabel status, and the annotation guideline. - :return: None - :raises LabellerrError: If the update fails. - """ - unique_id = str(uuid.uuid4()) - - url = f"{constants.BASE_URL}/annotations/create_template?data_type={data_type}&client_id={client_id}&uuid={unique_id}" - - guide_payload = json.dumps( - {"templateName": template_name, "questions": questions} - ) - - headers = client_utils.build_headers( - api_key=self.api_key, - api_secret=self.api_secret, - client_id=client_id, - extra_headers={"content-type": "application/json"}, - ) - - try: - response_data = client_utils.request( - "POST", url, headers=headers, data=guide_payload, request_id=unique_id - ) - return response_data["response"]["template_id"] - except requests.exceptions.RequestException as e: - logging.error(f"Failed to update project annotation guideline: {str(e)}") - raise - - def validate_rotation_config(self, rotation_config): - """ - Validates a rotation configuration. - - :param rotation_config: A dictionary containing the configuration for the rotations. - :raises LabellerrError: If the configuration is invalid. - """ - client_utils.validate_rotation_config(rotation_config) - - def create_dataset( - self, - dataset_config, - files_to_upload=None, - folder_to_upload=None, - connector_config=None, - ): - """ - Creates a dataset with support for multiple data types and connectors. - - :param dataset_config: A dictionary containing the configuration for the dataset. - Required fields: client_id, dataset_name, data_type - Optional fields: dataset_description, connector_type - :param files_to_upload: List of file paths to upload (for local connector) - :param folder_to_upload: Path to folder to upload (for local connector) - :param connector_config: Configuration for cloud connectors (GCP/AWS) - :return: A dictionary containing the response status and the ID of the created dataset. - """ - - try: - # Validate required fields - required_fields = ["client_id", "dataset_name", "data_type"] - for field in required_fields: - if field not in dataset_config: - raise LabellerrError( - f"Required field '{field}' missing in dataset_config" - ) - - # Validate data_type - if dataset_config.get("data_type") not in constants.DATA_TYPES: - raise LabellerrError( - f"Invalid data_type. Must be one of {constants.DATA_TYPES}" - ) - - connector_type = dataset_config.get("connector_type", "local") - connection_id = None - path = connector_type - - # Handle different connector types - if connector_type == "local": - if files_to_upload is not None: - try: - connection_id = self.client.upload_files( - client_id=dataset_config["client_id"], - files_list=files_to_upload, - ) - except Exception as e: - raise LabellerrError( - f"Failed to upload files to dataset: {str(e)}" - ) - - elif folder_to_upload is not None: - try: - result = self.upload_folder_files_to_dataset( - { - "client_id": dataset_config["client_id"], - "folder_path": folder_to_upload, - "data_type": dataset_config["data_type"], - } - ) - connection_id = result["connection_id"] - except Exception as e: - raise LabellerrError( - f"Failed to upload folder files to dataset: {str(e)}" - ) - elif connector_config is None: - # Create empty dataset for local connector - connection_id = None - - elif connector_type in ["gcp", "aws"]: - if connector_config is None: - raise LabellerrError( - f"connector_config is required for {connector_type} connector" - ) - - try: - connection_id = self.client._setup_cloud_connector( - connector_type, dataset_config["client_id"], connector_config - ) - except Exception as e: - raise LabellerrError( - f"Failed to setup {connector_type} connector: {str(e)}" - ) - else: - raise LabellerrError(f"Unsupported connector type: {connector_type}") - - unique_id = str(uuid.uuid4()) - url = f"{constants.BASE_URL}/datasets/create?client_id={dataset_config['client_id']}&uuid={unique_id}" - headers = client_utils.build_headers( - api_key=self.api_key, - api_secret=self.api_secret, - client_id=dataset_config["client_id"], - extra_headers={"content-type": "application/json"}, - ) - - payload = json.dumps( - { - "dataset_name": dataset_config["dataset_name"], - "dataset_description": dataset_config.get( - "dataset_description", "" - ), - "data_type": dataset_config["data_type"], - "connection_id": connection_id, - "path": path, - "client_id": dataset_config["client_id"], - "connector_type": connector_type, - } - ) - response_data = client_utils.request( - "POST", url, headers=headers, data=payload, request_id=unique_id - ) - dataset_id = response_data["response"]["dataset_id"] - - return {"response": "success", "dataset_id": dataset_id} - - except LabellerrError as e: - logging.error(f"Failed to create dataset: {e}") - raise - - def delete_dataset(self, client_id, dataset_id): - """ - Deletes a dataset from the system. - - :param client_id: The ID of the client - :param dataset_id: The ID of the dataset to delete - :return: Dictionary containing deletion status - :raises LabellerrError: If the deletion fails - """ - # Validate parameters using Pydantic - params = schemas.DeleteDatasetParams(client_id=client_id, dataset_id=dataset_id) - unique_id = str(uuid.uuid4()) - url = f"{constants.BASE_URL}/datasets/{params.dataset_id}/delete?client_id={params.client_id}&uuid={unique_id}" - headers = client_utils.build_headers( - api_key=self.api_key, - api_secret=self.api_secret, - client_id=params.client_id, - extra_headers={"content-type": "application/json"}, - ) - - return client_utils.request( - "DELETE", url, headers=headers, request_id=unique_id - ) - - def upload_folder_files_to_dataset(self, data_config): - """ - Uploads local files from a folder to a dataset using parallel processing. - - :param data_config: A dictionary containing the configuration for the data. - :return: A dictionary containing the response status and the list of successfully uploaded files. - :raises LabellerrError: If there are issues with file limits, permissions, or upload process - """ - try: - # Validate required fields in data_config - required_fields = ["client_id", "folder_path", "data_type"] - missing_fields = [ - field for field in required_fields if field not in data_config - ] - if missing_fields: - raise LabellerrError( - f"Missing required fields in data_config: {', '.join(missing_fields)}" - ) - - # Validate folder path exists and is accessible - if not os.path.exists(data_config["folder_path"]): - raise LabellerrError( - f"Folder path does not exist: {data_config['folder_path']}" - ) - if not os.path.isdir(data_config["folder_path"]): - raise LabellerrError( - f"Path is not a directory: {data_config['folder_path']}" - ) - if not os.access(data_config["folder_path"], os.R_OK): - raise LabellerrError( - f"No read permission for folder: {data_config['folder_path']}" - ) - - success_queue = [] - fail_queue = [] - - try: - # Get files from folder - total_file_count, total_file_volumn, filenames = ( - self.client.get_total_folder_file_count_and_total_size( - data_config["folder_path"], data_config["data_type"] - ) - ) - except Exception as e: - logging.error(f"Failed to analyze folder contents: {str(e)}") - raise - - # Check file limits - if total_file_count > constants.TOTAL_FILES_COUNT_LIMIT_PER_DATASET: - raise LabellerrError( - f"Total file count: {total_file_count} exceeds limit of {constants.TOTAL_FILES_COUNT_LIMIT_PER_DATASET} files" - ) - if total_file_volumn > constants.TOTAL_FILES_SIZE_LIMIT_PER_DATASET: - raise LabellerrError( - f"Total file size: {total_file_volumn/1024/1024:.1f}MB exceeds limit of {constants.TOTAL_FILES_SIZE_LIMIT_PER_DATASET/1024/1024:.1f}MB" - ) - - logging.info(f"Total file count: {total_file_count}") - logging.info(f"Total file size: {total_file_volumn/1024/1024:.1f} MB") - - # Use generator for memory-efficient batch creation - def create_batches(): - current_batch = [] - current_batch_size = 0 - - for file_path in filenames: - try: - file_size = os.path.getsize(file_path) - if ( - current_batch_size + file_size > constants.FILE_BATCH_SIZE - or len(current_batch) >= constants.FILE_BATCH_COUNT - ): - if current_batch: - yield current_batch - current_batch = [file_path] - current_batch_size = file_size - else: - current_batch.append(file_path) - current_batch_size += file_size - except OSError as e: - logging.error(f"Error accessing file {file_path}: {str(e)}") - fail_queue.append(file_path) - except Exception as e: - logging.error( - f"Unexpected error processing {file_path}: {str(e)}" - ) - fail_queue.append(file_path) - - if current_batch: - yield current_batch - - # Convert generator to list for ThreadPoolExecutor - batches = list(create_batches()) - - if not batches: - raise LabellerrError( - "No valid files found to upload in the specified folder" - ) - - logging.info(f"CPU count: {os.cpu_count()}, Batch Count: {len(batches)}") - - # Calculate optimal number of workers based on CPU count and batch count - max_workers = min( - os.cpu_count(), # Number of CPU cores - len(batches), # Number of batches - 20, - ) - connection_id = str(uuid.uuid4()) - # Process batches in parallel - with ThreadPoolExecutor(max_workers=max_workers) as executor: - future_to_batch = { - executor.submit( - self.__process_batch, - data_config["client_id"], - batch, - connection_id, - ): batch - for batch in batches - } - - for future in as_completed(future_to_batch): - batch = future_to_batch[future] - try: - result = future.result() - if ( - isinstance(result, dict) - and result.get("message") == "200: Success" - ): - success_queue.extend(batch) - else: - fail_queue.extend(batch) - except Exception as e: - logging.exception(e) - logging.error(f"Batch upload failed: {str(e)}") - fail_queue.extend(batch) - - if not success_queue and fail_queue: - raise LabellerrError( - "All file uploads failed. Check individual file errors above." - ) - - return { - "connection_id": connection_id, - "success": success_queue, - "fail": fail_queue, - } - - except LabellerrError: - raise - except Exception as e: - logging.error(f"Failed to upload files: {str(e)}") - raise - - def __process_batch(self, client_id, files_list, connection_id=None): - """ - Processes a batch of files. - """ - # Prepare files for upload - files = {} - for file_path in files_list: - file_name = os.path.basename(file_path) - files[file_name] = file_path - - response = self.client.connect_local_files( - client_id, list(files.keys()), connection_id - ) - resumable_upload_links = response["response"]["resumable_upload_links"] - for file_name in resumable_upload_links.keys(): - gcs.upload_to_gcs_resumable( - resumable_upload_links[file_name], files[file_name] - ) - - return response - - def attach_dataset_to_project( - self, client_id, project_id, dataset_id=None, dataset_ids=None - ): - """ - Attaches one or more datasets to an existing project. - - :param client_id: The ID of the client - :param project_id: The ID of the project - :param dataset_id: The ID of a single dataset to attach (for backward compatibility) - :param dataset_ids: List of dataset IDs to attach (for batch operations) - :return: Dictionary containing attachment status - :raises LabellerrError: If the operation fails or if neither dataset_id nor dataset_ids is provided - """ - # Handle both single and batch operations - if dataset_id is None and dataset_ids is None: - raise LabellerrError("Either dataset_id or dataset_ids must be provided") - - if dataset_id is not None and dataset_ids is not None: - raise LabellerrError( - "Cannot provide both dataset_id and dataset_ids. Use dataset_ids for batch operations." - ) - - # Convert single dataset_id to list for uniform processing - if dataset_id is not None: - dataset_ids = [dataset_id] - - # Validate parameters using Pydantic for each dataset - validated_dataset_ids = [] - for ds_id in dataset_ids: - params = schemas.AttachDatasetParams( - client_id=client_id, project_id=project_id, dataset_id=ds_id - ) - validated_dataset_ids.append(str(params.dataset_id)) - - # Use the first params validation for client_id and project_id - params = schemas.AttachDatasetParams( - client_id=client_id, project_id=project_id, dataset_id=dataset_ids[0] - ) - - unique_id = str(uuid.uuid4()) - url = f"{constants.BASE_URL}/actions/jobs/add_datasets_to_project?project_id={params.project_id}&uuid={unique_id}&client_id={params.client_id}" - headers = client_utils.build_headers( - api_key=self.api_key, - api_secret=self.api_secret, - client_id=params.client_id, - extra_headers={"content-type": "application/json"}, - ) - - payload = json.dumps({"attached_datasets": validated_dataset_ids}) - - return client_utils.request( - "POST", url, headers=headers, data=payload, request_id=unique_id - ) - - def detach_dataset_from_project( - self, client_id, project_id, dataset_id=None, dataset_ids=None - ): - """ - Detaches one or more datasets from an existing project. - - :param client_id: The ID of the client - :param project_id: The ID of the project - :param dataset_id: The ID of a single dataset to detach (for backward compatibility) - :param dataset_ids: List of dataset IDs to detach (for batch operations) - :return: Dictionary containing detachment status - :raises LabellerrError: If the operation fails or if neither dataset_id nor dataset_ids is provided - """ - # Handle both single and batch operations - if dataset_id is None and dataset_ids is None: - raise LabellerrError("Either dataset_id or dataset_ids must be provided") - - if dataset_id is not None and dataset_ids is not None: - raise LabellerrError( - "Cannot provide both dataset_id and dataset_ids. Use dataset_ids for batch operations." - ) - - # Convert single dataset_id to list for uniform processing - if dataset_id is not None: - dataset_ids = [dataset_id] - - # Validate parameters using Pydantic for each dataset - validated_dataset_ids = [] - for ds_id in dataset_ids: - params = schemas.DetachDatasetParams( - client_id=client_id, project_id=project_id, dataset_id=ds_id - ) - validated_dataset_ids.append(str(params.dataset_id)) - - # Use the first params validation for client_id and project_id - params = schemas.DetachDatasetParams( - client_id=client_id, project_id=project_id, dataset_id=dataset_ids[0] - ) - - unique_id = str(uuid.uuid4()) - url = f"{constants.BASE_URL}/actions/jobs/delete_datasets_from_project?project_id={params.project_id}&uuid={unique_id}" - headers = client_utils.build_headers( - api_key=self.api_key, - api_secret=self.api_secret, - client_id=params.client_id, - extra_headers={"content-type": "application/json"}, - ) - - payload = json.dumps({"attached_datasets": validated_dataset_ids}) - - return client_utils.request( - "POST", url, headers=headers, data=payload, request_id=unique_id - ) - - @validate_params(client_id=str, datatype=str, project_id=str, scope=str) - def get_all_datasets( - self, client_id: str, datatype: str, project_id: str, scope: str - ): - """ - Retrieves datasets by parameters. - - :param client_id: The ID of the client. - :param datatype: The type of data for the dataset. - :param project_id: The ID of the project. - :param scope: The permission scope for the dataset. - :return: The dataset list as JSON. - """ - # Validate parameters using Pydantic - params = schemas.GetAllDatasetParams( - client_id=client_id, - datatype=datatype, - project_id=project_id, - scope=scope, - ) - unique_id = str(uuid.uuid4()) - url = ( - f"{constants.BASE_URL}/datasets/list?client_id={params.client_id}&data_type={params.datatype}&permission_level={params.scope}" - f"&project_id={params.project_id}&uuid={unique_id}" - ) - headers = client_utils.build_headers( - api_key=self.api_key, - api_secret=self.api_secret, - client_id=params.client_id, - extra_headers={"content-type": "application/json"}, - ) - - return client_utils.request("GET", url, headers=headers, request_id=unique_id) diff --git a/labellerr/core/datasets/utils.py b/labellerr/core/datasets/utils.py index a8fe268..cf1737d 100644 --- a/labellerr/core/datasets/utils.py +++ b/labellerr/core/datasets/utils.py @@ -2,14 +2,11 @@ import os import uuid from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import TYPE_CHECKING, List, Union +from typing import List, Union -from .. import client_utils, constants, gcs, schemas +from .. import client_utils, constants, gcs from ..exceptions import LabellerrError -from ..utils import validate_params - -if TYPE_CHECKING: - from ..client import LabellerrClient +from ..client import LabellerrClient def get_total_folder_file_count_and_total_size(folder_path, data_type): @@ -110,7 +107,6 @@ def connect_local_files( return client_utils.request("POST", url, headers=headers, json=body) -@validate_params(client_id=str, files_list=(str, list)) def upload_files( client: "LabellerrClient", client_id: str, files_list: Union[str, List[str]] ): @@ -122,23 +118,12 @@ def upload_files( :return: The connection ID from the API. :raises LabellerrError: If the upload fails. """ - # Validate parameters using Pydantic - params = schemas.UploadFilesParams(client_id=client_id, files_list=files_list) - try: - # Use validated files_list from Pydantic - files_list = params.files_list + if len(files_list) == 0: + raise LabellerrError("No files to upload") - if len(files_list) == 0: - raise LabellerrError("No files to upload") - - response = __process_batch(client, client_id, files_list) - connection_id = response["response"]["temporary_connection_id"] - return connection_id - except LabellerrError: - raise - except Exception as e: - logging.error(f"Failed to upload files: {str(e)}") - raise + response = __process_batch(client, client_id, files_list) + connection_id = response["response"]["temporary_connection_id"] + return connection_id def __process_batch( diff --git a/labellerr/core/exceptions/__init__.py b/labellerr/core/exceptions/__init__.py index ddb55c6..7b9e511 100644 --- a/labellerr/core/exceptions/__init__.py +++ b/labellerr/core/exceptions/__init__.py @@ -25,3 +25,9 @@ class InvalidDatasetIDError(Exception): class InvalidConnectionError(Exception): pass + + +class InvalidAnnotationTemplateError(Exception): + """Custom exception for invalid annotation template errors.""" + + pass diff --git a/labellerr/core/files/video_file.py b/labellerr/core/files/video_file.py index 7e544d5..2951882 100644 --- a/labellerr/core/files/video_file.py +++ b/labellerr/core/files/video_file.py @@ -1,10 +1,11 @@ +import json import os import shutil import subprocess import uuid from concurrent.futures import ThreadPoolExecutor, as_completed from threading import Lock -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Dict, List, Optional import requests @@ -25,7 +26,7 @@ def __init__( client: "LabellerrClient", file_id: str, project_id: str, - dataset_id: str | None = None, + dataset_id: Optional[str] = None, **kwargs, ): super().__init__(client, file_id, project_id, dataset_id=dataset_id, **kwargs) @@ -35,7 +36,7 @@ def total_frames(self): """Get total number of frames in the video.""" return self.metadata.get("total_frames", 0) - def get_frames(self, frame_start: int = 0, frame_end: int | None = None): + def get_frames(self, frame_start: int = 0, frame_end: Optional[int] = None): """ Retrieve video frames data from Labellerr API. @@ -102,7 +103,10 @@ def _download_single_frame(self, frame_number, frame_url, save_path, print_lock) return False, frame_number, error_info def download_frames( - self, frames_data: dict, output_folder: str | None = None, max_workers: int = 30 + self, + frames_data: dict, + output_folder: Optional[str] = None, + max_workers: int = 30, ): """ Download video frames from URLs to a local folder using multithreading. @@ -189,7 +193,7 @@ def create_video( frames_folder: str, framerate: int = 30, pattern: str = "%d.jpg", - output_file: str | None = None, + output_file: Optional[str] = None, ): """ Join frames into a video using ffmpeg. @@ -330,5 +334,89 @@ def download_create_video_auto_cleanup( raise LabellerrError(f"Failed in video processing: {str(e)}") + def upload_pre_annotations( + self, + annotations: List[Dict], + annotation_format: str = "video_json", + conf_bucket: str = None, + ): + """ + Upload pre-annotations for this video file. + + This is a convenience method that creates a properly formatted annotation file + for a single video and uploads it using the project's upload_preannotation method. + + :param annotations: List of question annotations in format: + [ + { + "question_name": "Label name", + "question_type": "BoundingBox" or "polygon", + "answer": [ + { + "frames": { + "0": { + "frame": 0, + "answer": { + "xmin": 100, "ymin": 100, "xmax": 300, "ymax": 300, "rotation": 0 + }, + "timestamp": 0.0 + }, + "25": { + "frame": 25, + "answer": { + "xmin": 150, "ymin": 150, "xmax": 350, "ymax": 350, "rotation": 0 + }, + "timestamp": 1.0 + } + } + } + ] + } + ] + + For polygon annotations, use answer format: + "answer": [{"x": 0, "y": 600}, {"x": 1920, "y": 600}, ...] + + :param annotation_format: Format of annotations (default: "video_json") + :param conf_bucket: Optional confidence bucket ("low", "medium", "high") + :return: Response with job status + :raises LabellerrError: If upload fails + """ + try: + # Get project instance + from ..projects.base import LabellerrProject + + project = LabellerrProject(self.client, self.project_id) + + # Create properly formatted video answer data + video_data = [ + { + "file_name": self.metadata.get("file_name", f"{self.file_id}.mp4"), + "annotations": annotations, + } + ] + + # Create temporary file with the annotations + temp_file_path = f"/tmp/{self.file_id}_preannotations.json" + with open(temp_file_path, "w") as f: + json.dump(video_data, f, indent=2) + + # Upload using project method + result = project.upload_preannotations( + annotation_format=annotation_format, + annotation_file=temp_file_path, + conf_bucket=conf_bucket, + _async=True, + ) + + # Clean up temp file + if os.path.exists(temp_file_path): + os.remove(temp_file_path) + + return result + + except Exception as e: + raise LabellerrError(f"Failed to upload video pre-annotations: {str(e)}") + LabellerrFileMeta._register("video", LabellerrVideoFile) diff --git a/labellerr/core/projects/__init__.py b/labellerr/core/projects/__init__.py index 60a41ce..67eac7b 100644 --- a/labellerr/core/projects/__init__.py +++ b/labellerr/core/projects/__init__.py @@ -6,8 +6,8 @@ from labellerr import LabellerrClient -from .. import client_utils, constants, schemas, utils -from ..datasets import LabellerrDataset, create_dataset +from .. import constants, schemas, utils +from ..datasets import LabellerrDataset, create_dataset_from_local from ..exceptions import LabellerrError from .audio_project import AudioProject as LabellerrAudioProject from .base import LabellerrProject @@ -17,11 +17,11 @@ from .video_project import VideoProject as LabellerrVideoProject __all__ = [ - "LabellerrImageProject", - "LabellerrVideoProject", "LabellerrProject", - "LabellerrDocumentProject", "LabellerrAudioProject", + "LabellerrDocumentProject", + "LabellerrImageProject", + "LabellerrVideoProject", ] @@ -32,6 +32,17 @@ def create_project(client: "LabellerrClient", payload: dict): """ try: + # Validate client_id + if "client_id" not in payload: + raise LabellerrError("Required parameter client_id is missing") + + # Validate client_id is a non-empty string + if ( + not isinstance(payload.get("client_id"), str) + or not payload["client_id"].strip() + ): + raise LabellerrError("client_id must be a non-empty string") + # validate all the parameters required_params = [ "data_type", @@ -44,6 +55,10 @@ def create_project(client: "LabellerrClient", payload: dict): if param not in payload: raise LabellerrError(f"Required parameter {param} is missing") + # Check for dataset_name when creating new dataset + if "datasets" not in payload and "dataset_name" not in payload: + raise LabellerrError("Required parameter dataset_name is missing") + # Validate created_by email format created_by = payload.get("created_by") if ( @@ -97,7 +112,7 @@ def create_project(client: "LabellerrClient", payload: dict): logging.info("Validating existing datasets . . .") for dataset_id in datasets: try: - dataset = LabellerrDataset(client, dataset_id) + dataset = LabellerrDataset(client, dataset_id) # type: ignore[abstract] if dataset.files_count <= 0: raise LabellerrError(f"Dataset {dataset_id} has no files") except Exception as e: @@ -147,7 +162,7 @@ def create_project(client: "LabellerrClient", payload: dict): logging.info("Creating dataset . . .") - dataset = create_dataset( + dataset = create_dataset_from_local( client, schemas.DatasetConfig( client_id=client.client_id, @@ -256,18 +271,14 @@ def __create_project_api_call( "created_by": params.created_by, } ) - headers = client_utils.build_headers( - api_key=client.api_key, - api_secret=client.api_secret, - client_id=params.client_id, - extra_headers={ - "Origin": constants.ALLOWED_ORIGINS, - "Content-Type": "application/json", - }, - ) return client.make_request( - "POST", url, headers=headers, data=payload, request_id=unique_id + "POST", + url, + client_id=params.client_id, + extra_headers={"content-type": "application/json"}, + request_id=unique_id, + data=payload, ) @@ -291,3 +302,25 @@ def create_annotation_guideline( except requests.exceptions.RequestException as e: logging.error(f"Failed to update project annotation guideline: {str(e)}") raise + + +def list_projects(client: "LabellerrClient"): + """ + Retrieves a list of projects associated with a client ID. + + :param client: The client instance. + :return: A list of LabellerrProject objects. + """ + unique_id = str(uuid.uuid4()) + url = f"{constants.BASE_URL}/project_drafts/projects/detailed_list?client_id={client.client_id}&uuid={unique_id}" + + response = client.make_request( + "GET", + url, + extra_headers={"content-type": "application/json"}, + request_id=unique_id, + ) + return [ + LabellerrProject(client, project_id=project["project_id"]) + for project in response["response"]["projects"] + ] diff --git a/labellerr/core/projects/annotation_guide.py b/labellerr/core/projects/annotation_guide.py new file mode 100644 index 0000000..e69de29 diff --git a/labellerr/core/projects/base.py b/labellerr/core/projects/base.py index 9af1c7a..14e05de 100644 --- a/labellerr/core/projects/base.py +++ b/labellerr/core/projects/base.py @@ -6,17 +6,16 @@ import os import uuid from abc import ABCMeta -from typing import TYPE_CHECKING, Dict +from typing import Dict import requests -from .. import client_utils, constants, gcs, schemas +from .. import client_utils, constants, schemas from ..exceptions import InvalidProjectError, LabellerrError -from ..utils import poll +from .utils import poll from ..exports import Export -if TYPE_CHECKING: - from ..client import LabellerrClient +from ..client import LabellerrClient class LabellerrProjectMeta(ABCMeta): @@ -59,8 +58,6 @@ def __call__(cls, client, project_id, **kwargs): if project_data is None: raise InvalidProjectError(f"Project not found: {project_id}") data_type = project_data.get("data_type") - if data_type not in constants.DATA_TYPES: - raise InvalidProjectError(f"Data type not supported: {data_type}") project_class = cls._registry.get(data_type) if project_class is None: @@ -75,44 +72,31 @@ class LabellerrProject(metaclass=LabellerrProjectMeta): def __init__(self, client: "LabellerrClient", project_id: str, **kwargs): self.client = client self.project_id = project_id - self.project_data = kwargs["project_data"] + self.__project_data = kwargs["project_data"] @property - def data_type(self): - return self.project_data.get("data_type") + def status_code(self): + return self.__project_data.get("status_code", 501) # if not found, return 501 @property - def attached_datasets(self): - return self.project_data.get("attached_datasets") + def annotation_template_id(self): + return self.__project_data.get("annotation_template_id") - def get_direct_upload_url( - self, file_name: str, client_id: str, purpose: str = "pre-annotations" - ) -> str: - """ - Get a direct upload URL for uploading files to GCS. + @property + def created_by(self): + return self.__project_data.get("created_by") - :param file_name: Name of the file to upload - :param client_id: Client ID - :param purpose: Purpose of the upload (default: "pre-annotations") - :return: Direct upload URL - """ - url = f"{constants.BASE_URL}/connectors/direct-upload-url" - params = { # noqa: F841 - "client_id": client_id, - "purpose": purpose, - "file_name": file_name, - } + @property + def created_at(self): + return self.__project_data.get("created_at") - try: - response_data = self.client.make_request( - "GET", - url, - extra_headers={"Origin": constants.ALLOWED_ORIGINS}, - ) - return response_data["response"] - except Exception as e: - logging.error(f"Error getting direct upload url: {e}") - raise LabellerrError(f"Failed to get direct upload URL: {str(e)}") + @property + def data_type(self): + return self.__project_data.get("data_type") + + @property + def attached_datasets(self): + return self.__project_data.get("attached_datasets") def detach_dataset_from_project(self, dataset_id=None, dataset_ids=None): """ @@ -247,227 +231,6 @@ def update_rotation_count(self, rotation_config): logging.error(f"Project rotation update config failed: {e}") raise - @staticmethod - def list_all_projects(client: "LabellerrClient"): - """ - Retrieves a list of projects associated with a client ID. - - :param client: The client instance. - :return: A dictionary containing the list of projects. - :raises LabellerrError: If the retrieval fails. - """ - try: - unique_id = str(uuid.uuid4()) - url = f"{constants.BASE_URL}/project_drafts/projects/detailed_list?client_id={client.client_id}&uuid={unique_id}" - - return client.make_request( - "GET", - url, - extra_headers={"content-type": "application/json"}, - request_id=unique_id, - ) - except Exception as e: - logging.error(f"Failed to retrieve projects: {str(e)}") - raise - - def _upload_preannotation_sync( - self, - project_id, - client_id, - annotation_format, - annotation_file, - conf_bucket=None, - ): - """ - Synchronous implementation of preannotation upload. - - :param project_id: The ID of the project. - :param client_id: The ID of the client. - :param annotation_format: The format of the preannotation data. - :param annotation_file: The file path of the preannotation data. - :param conf_bucket: Confidence bucket [low, medium, high] - :return: The response from the API. - :raises LabellerrError: If the upload fails. - """ - try: - # validate all the parameters - required_params = { - "project_id": project_id, - "client_id": client_id, - "annotation_format": annotation_format, - "annotation_file": annotation_file, - } - client_utils.validate_required_params( - required_params, list(required_params.keys()) - ) - client_utils.validate_annotation_format(annotation_format, annotation_file) - - request_uuid = str(uuid.uuid4()) - url = f"{constants.BASE_URL}/actions/upload_answers?project_id={project_id}&answer_format={annotation_format}&client_id={client_id}&uuid={request_uuid}" - if conf_bucket: - assert conf_bucket in [ - "low", - "medium", - "high", - ], "Invalid confidence bucket value. Must be one of [low, medium, high]" - url += f"&conf_bucket={conf_bucket}" - file_name = client_utils.validate_file_exists(annotation_file) - # get the direct upload url - gcs_path = f"{project_id}/{annotation_format}-{file_name}" - logging.info("Uploading your file to Labellerr. Please wait...") - direct_upload_url = self.get_direct_upload_url(gcs_path, client_id) - # Now let's wait for the file to be uploaded to the gcs - gcs.upload_to_gcs_direct(direct_upload_url, annotation_file) - payload = {} - url += "&gcs_path=" + gcs_path - - response = self.client.make_request( - "POST", - url, - extra_headers={"email_id": self.client.api_key}, - request_id=request_uuid, - handle_response=False, - data=payload, - ) - response_data = self.client.handle_upload_response(response, request_uuid) - - # read job_id from the response - job_id = response_data["response"]["job_id"] - self.client_id = client_id - - logging.info(f"Preannotation upload successful. Job ID: {job_id}") - - # Use max_retries=10 with 5-second intervals = 50 seconds max (fits within typical test timeouts) - future = self.preannotation_job_status_async(project_id, job_id) - return future.result() - except Exception as e: - logging.error(f"Failed to upload preannotation: {str(e)}") - raise - - def upload_preannotation_async( - self, annotation_format, annotation_file, conf_bucket=None - ): - """ - Asynchronously uploads preannotation data to a project. - - :param annotation_format: The format of the preannotation data. - :param annotation_file: The file path of the preannotation data. - :param conf_bucket: Confidence bucket [low, medium, high] - :return: A Future object that will contain the response from the API. - :raises LabellerrError: If the upload fails. - - """ - - def upload_and_monitor(): - try: - # validate all the parameters - required_params = [ - "annotation_format", - "annotation_file", - ] - for param in required_params: - if param not in locals(): - raise LabellerrError(f"Required parameter {param} is missing") - - if annotation_format not in constants.ANNOTATION_FORMAT: - raise LabellerrError( - f"Invalid annotation_format. Must be one of {constants.ANNOTATION_FORMAT}" - ) - - request_uuid = str(uuid.uuid4()) - url = ( - f"{constants.BASE_URL}/actions/upload_answers?" - f"project_id={self.project_id}&answer_format={annotation_format}&client_id={self.client.client_id}&uuid={request_uuid}" - ) - if conf_bucket: - assert conf_bucket in [ - "low", - "medium", - "high", - ], "Invalid confidence bucket value. Must be one of [low, medium, high]" - url += f"&conf_bucket={conf_bucket}" - # validate if the file exist then extract file name from the path - if os.path.exists(annotation_file): - file_name = os.path.basename(annotation_file) - else: - raise LabellerrError("File not found") - - # Check if the file extension is .json when annotation_format is coco_json - if annotation_format == "coco_json": - file_extension = os.path.splitext(annotation_file)[1].lower() - if file_extension != ".json": - raise LabellerrError( - "For coco_json annotation format, the file must have a .json extension" - ) - # get the direct upload url - gcs_path = f"{self.project_id}/{annotation_format}-{file_name}" - logging.info("Uploading your file to Labellerr. Please wait...") - direct_upload_url = self.get_direct_upload_url( - gcs_path, self.client.client_id - ) - # Now let's wait for the file to be uploaded to the gcs - gcs.upload_to_gcs_direct(direct_upload_url, annotation_file) - payload = {} - url += "&gcs_path=" + gcs_path - - response = self.client.make_request( - "POST", - url, - extra_headers={"email_id": self.client.api_key}, - request_id=request_uuid, - handle_response=False, - data=payload, - ) - response_data = self.client.handle_upload_response( - response, request_uuid - ) - - # read job_id from the response - job_id = response_data["response"]["job_id"] - - logging.info(f"Pre annotation upload successful. Job ID: {job_id}") - - # Now monitor the status - status_url = f"{constants.BASE_URL}/actions/upload_answers_status?project_id={self.project_id}&job_id={job_id}&client_id={self.client.client_id}" - - def check_job_status(): - status_data = self.client.make_request( - "GET", - status_url, - extra_headers={"Origin": constants.ALLOWED_ORIGINS}, - ) - logging.debug(f"Status data: {status_data}") - return status_data - - def is_job_completed(status_data): - return status_data.get("response", {}).get("status") == "completed" - - def on_success(status_data): - logging.info("Pre-annotation job completed.") - - def on_exception(e): - logging.error(f"Failed to get preannotation job status: {str(e)}") - raise LabellerrError( - f"Failed to get preannotation job status: {str(e)}" - ) - - result = poll( - function=check_job_status, - condition=is_job_completed, - interval=5.0, - on_success=on_success, - on_exception=on_exception, - ) - - return result - - except Exception as e: - logging.exception(f"Failed to upload preannotation: {str(e)}") - raise - - with concurrent.futures.ThreadPoolExecutor() as executor: - return executor.submit(upload_and_monitor) - def preannotation_job_status_async(self, job_id): """ Get the status of a preannotation job asynchronously with timeout protection. @@ -521,16 +284,19 @@ def on_exception(e): with concurrent.futures.ThreadPoolExecutor() as executor: return executor.submit(check_status) - def upload_preannotations( - self, annotation_format, annotation_file, conf_bucket=None + def __upload_preannotations( + self, + annotation_format, + annotation_file, + conf_bucket=None, ): """ - Uploads preannotation data to a project. + Uploads preannotation data to a project asynchronously. :param annotation_format: The format of the preannotation data. :param annotation_file: The file path of the preannotation data. :param conf_bucket: Confidence bucket [low, medium, high] - :return: The response from the API. + :return: concurrent.futures.Future - call .result() to wait for completion :raises LabellerrError: If the upload fails. """ try: @@ -582,13 +348,32 @@ def upload_preannotations( logging.info(f"Preannotation job started successfully. Job ID: {job_id}") - # Use max_retries=10 with 5-second intervals = 50 seconds max (fits within typical test timeouts) - future = self.preannotation_job_status_async(job_id) - return future.result() + return self.preannotation_job_status_async(job_id) except Exception as e: logging.error(f"Failed to upload preannotation: {str(e)}") raise + def upload_preannotations( + self, annotation_format, annotation_file, conf_bucket=None, _async=False + ): + """ + Backward compatibility method for upload_preannotations. + + :param annotation_format: The format of the preannotation data. + :param annotation_file: The file path of the preannotation data. + :param conf_bucket: Confidence bucket [low, medium, high] + :param _async: Whether to return a future object (True) or block until completion (False) + :return: The response from the API (blocks until completion). + :raises LabellerrError: If the upload fails. + """ + future = self.__upload_preannotations( + annotation_format, annotation_file, conf_bucket + ) + if _async: + return future + else: + return future.result() + def create_export(self, export_config: schemas.CreateExportParams): """ Creates an export with the given configuration. diff --git a/labellerr/core/projects/document_project.py b/labellerr/core/projects/document_project.py index 91dbb32..8885f94 100644 --- a/labellerr/core/projects/document_project.py +++ b/labellerr/core/projects/document_project.py @@ -1,11 +1,6 @@ -from typing import TYPE_CHECKING - from ..schemas import DatasetDataType from .base import LabellerrProject, LabellerrProjectMeta -if TYPE_CHECKING: - from ..client import LabellerrClient # noqa:F401 - class DocucmentProject(LabellerrProject): diff --git a/labellerr/core/projects/utils.py b/labellerr/core/projects/utils.py index 0060a3d..4490afa 100644 --- a/labellerr/core/projects/utils.py +++ b/labellerr/core/projects/utils.py @@ -1,7 +1,14 @@ from typing import Any, Dict from ..exceptions import LabellerrError -from ..utils import poll # noqa: F401 +from ..utils import poll + +from .. import constants +from ..client import LabellerrClient + +__all__ = [ + "poll", +] def validate_rotation_config(rotation_config: Dict[str, Any]) -> None: @@ -40,3 +47,32 @@ def validate_rotation_config(rotation_config: Dict[str, Any]) -> None: raise LabellerrError( "client_review_rotation_count must be 0 when annotation_rotation_count is greater than 1" ) + + +def get_direct_upload_url( + client: LabellerrClient, file_name: str, purpose: str = "pre-annotations" +) -> str: + """ + Get a direct upload URL for uploading files to GCS. + + :param file_name: Name of the file to upload + :param client: LabellerrClient instance + :param purpose: Purpose of the upload (default: "pre-annotations") + :return: Direct upload URL + """ + url = f"{constants.BASE_URL}/connectors/direct-upload-url" + params = { # noqa: F841 + "client_id": client.client_id, + "purpose": purpose, + "file_name": file_name, + } + + try: + response_data = client.make_request( + "GET", + url, + extra_headers={"Origin": constants.ALLOWED_ORIGINS}, + ) + return response_data["response"] + except Exception as e: + raise LabellerrError(f"Failed to get direct upload URL: {str(e)}") diff --git a/labellerr/core/projects/video_project.py b/labellerr/core/projects/video_project.py index cac694b..0b512eb 100644 --- a/labellerr/core/projects/video_project.py +++ b/labellerr/core/projects/video_project.py @@ -1,4 +1,5 @@ import uuid +from concurrent.futures import Future from typing import List from .. import constants @@ -82,5 +83,82 @@ def delete_keyframes(self, file_id: str, keyframes: List[int]): except Exception as e: raise LabellerrError(f"Failed to delete key frames: {str(e)}") + def upload_preannotation( + self, annotation_format: str, annotation_file: str, conf_bucket: str = None + ): + """ + Upload pre-annotations for video project. + + For video projects, use annotation_format="video_json" with a JSON file containing + video annotations in the format: + [ + { + "file_name": "video.mp4", + "annotations": [ + { + "question_name": "Label name", + "question_type": "BoundingBox" or "polygon", + "answer": [ + { + "frames": { + "0": { + "frame": 0, + "answer": { + "xmin": 100, "ymin": 100, "xmax": 300, "ymax": 300, "rotation": 0 + }, + "timestamp": 0.0 + }, + "25": { + "frame": 25, + "answer": { + "xmin": 150, "ymin": 150, "xmax": 350, "ymax": 350, "rotation": 0 + }, + "timestamp": 1.0 + } + } + } + ] + } + ] + } + ] + + For polygon annotations, use answer format: + "answer": [{"x": 0, "y": 600}, {"x": 1920, "y": 600}, ...] + + :param annotation_format: Format of annotations ("video_json", "coco_json", etc.) + :param annotation_file: Path to the annotation file + :param conf_bucket: Optional confidence bucket ("low", "medium", "high") + :return: Response with job status + :raises LabellerrError: If upload fails + """ + # Delegate to the base class synchronous upload (blocks until completion) + return self.upload_preannotations( + annotation_format, annotation_file, conf_bucket, _async=False + ) + + def upload_preannotation_async( + self, annotation_format: str, annotation_file: str, conf_bucket: str = None + ) -> Future: + """ + Asynchronously upload pre-annotations for video project and monitor the job status. + + This method returns immediately with a Future object. The actual upload and monitoring + happens in a background thread. Use future.result() to wait for completion. + + For video projects, use annotation_format="video_json" with a JSON file containing + video annotations. + + :param annotation_format: Format of annotations ("video_json", "coco_json", etc.) + :param annotation_file: Path to the annotation file + :param conf_bucket: Optional confidence bucket ("low", "medium", "high") + :return: Future object that will contain the response when complete + :raises LabellerrError: If upload fails + """ + # Delegate to the base class async upload (returns Future immediately) + return self.upload_preannotations( + annotation_format, annotation_file, conf_bucket, _async=True + ) + LabellerrProjectMeta._register(DatasetDataType.video, VideoProject) diff --git a/labellerr/core/schemas/__init__.py b/labellerr/core/schemas/__init__.py index dee02d0..9fa4967 100644 --- a/labellerr/core/schemas/__init__.py +++ b/labellerr/core/schemas/__init__.py @@ -18,14 +18,15 @@ from labellerr.core.schemas.base import DirPathStr, FilePathStr, NonEmptyStr # Connection schemas -from labellerr.core.schemas.connections import ( +from labellerr.core.schemas.connectors import ( AWSConnectionParams, - AWSConnectorConfig, DatasetDataType, DeleteConnectionParams, - GCPConnectorConfig, GCSConnectionParams, AWSConnectionTestParams, + ConnectionType, + ConnectorType, + GCSConnectionTestParams, ) # Dataset schemas @@ -83,9 +84,10 @@ "AWSConnectionParams", "AWSConnectionTestParams", "GCSConnectionParams", + "GCSConnectionTestParams", "DeleteConnectionParams", - "AWSConnectorConfig", - "GCPConnectorConfig", + "ConnectorType", + "ConnectionType", "DatasetDataType", # Dataset schemas "UploadFilesParams", diff --git a/labellerr/core/schemas/annotation_templates.py b/labellerr/core/schemas/annotation_templates.py new file mode 100644 index 0000000..7016015 --- /dev/null +++ b/labellerr/core/schemas/annotation_templates.py @@ -0,0 +1,42 @@ +from pydantic import BaseModel +from typing import List, Optional +from enum import Enum +from ..schemas import DatasetDataType + + +class QuestionType(str, Enum): + bounding_box = "BoundingBox" + polygon = "polygon" + polyline = "polyline" + dot = "dot" + input = "input" + radio = "radio" + boolean = "boolean" + select = "select" + dropdown = "dropdown" + stt = "stt" + imc = "imc" + + +class Option(BaseModel): + option_name: str + + +class AnnotationQuestion(BaseModel): + """Question structure for annotation templates.""" + + question_number: int + question: str + question_id: str + question_type: QuestionType + required: bool + options: Optional[List[Option]] = [] + color: Optional[str] = None + + +class CreateTemplateParams(BaseModel): + """Parameters for creating an annotation template.""" + + template_name: str + data_type: DatasetDataType + questions: List[AnnotationQuestion] diff --git a/labellerr/core/schemas/base.py b/labellerr/core/schemas/base.py index 39cde3e..789c68a 100644 --- a/labellerr/core/schemas/base.py +++ b/labellerr/core/schemas/base.py @@ -3,6 +3,17 @@ """ import os +from enum import Enum + + +class DatasetDataType(str, Enum): + """Enum for dataset data types.""" + + image = "image" + video = "video" + audio = "audio" + document = "document" + text = "text" class NonEmptyStr(str): diff --git a/labellerr/core/schemas/connections.py b/labellerr/core/schemas/connections.py deleted file mode 100644 index 741b460..0000000 --- a/labellerr/core/schemas/connections.py +++ /dev/null @@ -1,100 +0,0 @@ -""" -Schema models for connection operations (AWS, GCS, etc.). -""" - -import os -from enum import Enum -from typing import Literal, Optional - -from pydantic import BaseModel, Field, field_validator - - -class DatasetDataType(str, Enum): - """Enum for dataset data types.""" - - image = "image" - video = "video" - audio = "audio" - document = "document" - text = "text" - - -class ConnectionType(str, Enum): - """Enum for connection types.""" - - _IMPORT = "import" - _EXPORT = "export" - - -class AWSConnectionTestParams(BaseModel): - """Parameters for testing an AWS S3 connection.""" - - aws_access_key: str = Field(min_length=1) - aws_secrets_key: str = Field(min_length=1) - s3_path: str = Field(min_length=1) - connection_type: ConnectionType = ConnectionType._IMPORT - - -class AWSConnectionParams(AWSConnectionTestParams): - """Parameters for creating an AWS S3 connection.""" - - name: str = Field(min_length=1) - description: str - - -class GCSConnectionParams(BaseModel): - """Parameters for creating a GCS connection.""" - - client_id: str = Field(min_length=1) - gcs_cred_file: str - gcs_path: str = Field(min_length=1) - data_type: DatasetDataType - name: str = Field(min_length=1) - description: str - connection_type: str = "import" - credentials: str = "svc_account_json" - - @field_validator("gcs_cred_file") - @classmethod - def validate_gcs_cred_file(cls, v): - if not os.path.exists(v): - raise ValueError(f"GCS credential file not found: {v}") - return v - - -class DeleteConnectionParams(BaseModel): - """Parameters for deleting a connection.""" - - client_id: str = Field(min_length=1) - connection_id: str = Field(min_length=1) - - -class AWSConnectorConfig(BaseModel): - """Configuration for AWS S3 connector.""" - - aws_access_key: str = Field(min_length=1) - aws_secrets_key: str = Field(min_length=1) - s3_path: str = Field(min_length=1) - data_type: Literal["image", "video", "audio", "document", "text"] - name: Optional[str] = None - description: str = "Auto-created AWS connector" - connection_type: str = "import" - - -class GCPConnectorConfig(BaseModel): - """Configuration for GCP connector.""" - - gcs_cred_file: str = Field(min_length=1) - gcs_path: str = Field(min_length=1) - data_type: Literal["image", "video", "audio", "document", "text"] - name: Optional[str] = None - description: str = "Auto-created GCS connector" - connection_type: str = "import" - credentials: str = "svc_account_json" - - @field_validator("gcs_cred_file") - @classmethod - def validate_gcs_cred_file(cls, v): - if not os.path.exists(v): - raise ValueError(f"GCS credential file not found: {v}") - return v diff --git a/labellerr/core/schemas/connectors.py b/labellerr/core/schemas/connectors.py new file mode 100644 index 0000000..aa3eaf4 --- /dev/null +++ b/labellerr/core/schemas/connectors.py @@ -0,0 +1,71 @@ +""" +Schema models for connection operations (AWS, GCS, etc.). +""" + +import os +from enum import Enum +from typing import Optional +from .base import DatasetDataType +from pydantic import BaseModel, Field, field_validator + + +class ConnectionType(str, Enum): + """Enum for connection types.""" + + _IMPORT = "import" + _EXPORT = "export" + + +class ConnectorType(str, Enum): + """Enum for connector types.""" + + _S3 = "s3" + _GCS = "gcs" + _LOCAL = "local" + + +class GCSConnectionTestParams(BaseModel): + """Parameters for testing a GCS connection.""" + + svc_account_json: Optional[str] = Field(default=None, min_length=1) + path: str = Field(min_length=1) + connection_type: ConnectionType = ConnectionType._IMPORT + data_type: DatasetDataType + + @field_validator("svc_account_json") + @classmethod + def validate_svc_account_json(cls, v): + if v and not os.path.exists(v): + raise ValueError(f"GCS credential file not found: {v}") + return v + + +class GCSConnectionParams(GCSConnectionTestParams): + """Parameters for creating a GCS connection.""" + + name: str = Field(min_length=1) + description: str + + +class AWSConnectionTestParams(BaseModel): + """Parameters for testing an AWS S3 connection.""" + + aws_access_key: str = Field(min_length=1) + aws_secrets_key: str = Field(min_length=1) + path: str = Field(min_length=1) + connection_type: ConnectionType = ConnectionType._IMPORT + data_type: DatasetDataType + + +class AWSConnectionParams(AWSConnectionTestParams): + """Parameters for creating an AWS S3 connection.""" + + name: str = Field(min_length=1) + description: str + + +class DeleteConnectionParams(BaseModel): + """Parameters for deleting a connection.""" + + client_id: str = Field(min_length=1) + connection_id: str = Field(min_length=1) diff --git a/labellerr/core/schemas/datasets.py b/labellerr/core/schemas/datasets.py index 6901a3c..2ada187 100644 --- a/labellerr/core/schemas/datasets.py +++ b/labellerr/core/schemas/datasets.py @@ -7,6 +7,7 @@ from typing import List, Literal from uuid import UUID +from .base import DatasetDataType from pydantic import BaseModel, Field, field_validator @@ -107,6 +108,6 @@ class DatasetConfig(BaseModel): """Configuration for creating a dataset.""" dataset_name: str = Field(min_length=1) - data_type: Literal["image", "video", "audio", "document", "text"] + data_type: DatasetDataType dataset_description: str = "" - connector_type: Literal["local", "aws", "gcp"] = "local" + multimodal_indexing: bool = False diff --git a/labellerr/core/schemas/projects.py b/labellerr/core/schemas/projects.py index 6e01333..830853b 100644 --- a/labellerr/core/schemas/projects.py +++ b/labellerr/core/schemas/projects.py @@ -4,7 +4,9 @@ from typing import Any, Dict, List, Literal, Optional -from pydantic import BaseModel, Field, field_validator +from pydantic import BaseModel, Field + +from .base import DatasetDataType class RotationConfig(BaseModel): @@ -38,23 +40,12 @@ class CreateProjectParams(BaseModel): """Parameters for creating a project.""" project_name: str = Field(min_length=1) - data_type: Literal["image", "video", "audio", "document", "text"] - client_id: str = Field(min_length=1) - attached_datasets: List[str] = Field(min_length=1) - annotation_template_id: str + data_type: DatasetDataType rotations: RotationConfig use_ai: bool = False - created_by: Optional[str] = None - - @field_validator("attached_datasets") - @classmethod - def validate_attached_datasets(cls, v): - if not v: - raise ValueError("must contain at least one dataset ID") - for i, dataset_id in enumerate(v): - if not isinstance(dataset_id, str) or not dataset_id.strip(): - raise ValueError(f"dataset_id at index {i} must be a non-empty string") - return v + created_by: Optional[str] = Field( + None, pattern=r"^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$" + ) class CreateTemplateParams(BaseModel): diff --git a/tests/integration/.gitignore b/tests/integration/.gitignore deleted file mode 100644 index 4134b79..0000000 --- a/tests/integration/.gitignore +++ /dev/null @@ -1,3 +0,0 @@ -__pycache__ -.env -.venv diff --git a/tests/integration/test_labellerr_integration.py b/tests/integration/test_labellerr_integration.py index f19ff48..2629803 100644 --- a/tests/integration/test_labellerr_integration.py +++ b/tests/integration/test_labellerr_integration.py @@ -137,12 +137,11 @@ def test_pre_annotation_upload_coco_json( annotation_file = temp_json_file(sample_annotation_data["coco_json"]) try: - result = project._upload_preannotation_sync( - project_id=test_project_ids["project_id"], - client_id=test_credentials["client_id"], + future = project.upload_preannotation( annotation_format="coco_json", annotation_file=annotation_file, ) + result = future.result() assert isinstance(result, dict) assert "response" in result @@ -183,12 +182,11 @@ def timeout_handler(signum, frame): signal.alarm(60) try: - result = project._upload_preannotation_sync( - project_id=test_project_ids["project_id"], - client_id=test_credentials["client_id"], + future = project.upload_preannotation( annotation_format="json", annotation_file=annotation_file, ) + result = future.result() assert isinstance(result, dict) @@ -230,12 +228,11 @@ def test_pre_annotation_invalid_format( project = LabellerrProject(integration_client, test_project_ids["project_id"]) with pytest.raises(LabellerrError) as exc_info: - project._upload_preannotation_sync( - project_id=test_project_ids["project_id"], - client_id=test_credentials["client_id"], + future = project.upload_preannotation( annotation_format=invalid_format, annotation_file="test.json", ) + future.result() assert expected_error in str(exc_info.value) @@ -400,7 +397,7 @@ def test_aws_connection_lifecycle(self, integration_client, test_credentials): client_id=test_credentials["client_id"], aws_access_key=aws_secret.get("access_key"), aws_secrets_key=aws_secret.get("secret_key"), - s3_path=aws_secret.get("s3_path"), + path=aws_secret.get("s3_path"), data_type=DatasetDataType.image, name=connection_name, description="Test AWS connection", diff --git a/tests/integration/test_sync_datasets.py b/tests/integration/test_sync_datasets.py index 902b76e..bdd7db2 100644 --- a/tests/integration/test_sync_datasets.py +++ b/tests/integration/test_sync_datasets.py @@ -89,7 +89,7 @@ def test_sync_datasets_aws(self): print(f"Data Type: {self.data_type}") print(f"Email ID: {self.email_id}") - response = datasets.sync_datasets( + response = datasets.sync_with_connection( client_id=self.client_id, project_id=self.project_id, dataset_id=self.aws_dataset_id, @@ -134,7 +134,7 @@ def test_sync_datasets_gcs(self): print(f"Data Type: {self.data_type}") print(f"Email ID: {self.email_id}") - response = datasets.sync_datasets( + response = datasets.sync_with_connection( client_id=self.client_id, project_id=self.project_id, dataset_id=self.gcs_dataset_id, @@ -169,7 +169,7 @@ def test_sync_datasets_with_multiple_data_types(self): print(f"\n Testing with data_type: {data_type}") try: - response = datasets.sync_datasets( + response = datasets.sync_with_connection( client_id=self.client_id, project_id=self.project_id, dataset_id=self.aws_dataset_id, @@ -194,7 +194,7 @@ def test_sync_datasets_invalid_connection_id(self): print("=" * 60) with self.assertRaises((LabellerrError, Exception)) as context: - datasets.sync_datasets( + datasets.sync_with_connection( client_id=self.client_id, project_id=self.project_id, dataset_id=self.aws_dataset_id, @@ -214,7 +214,7 @@ def test_sync_datasets_invalid_dataset_id(self): print("=" * 60) with self.assertRaises((LabellerrError, Exception)) as context: - datasets.sync_datasets( + datasets.sync_with_connection( client_id=self.client_id, project_id=self.project_id, dataset_id="00000000-0000-0000-0000-000000000000", diff --git a/tests/integration/test_video_preannotation_integration.py b/tests/integration/test_video_preannotation_integration.py new file mode 100644 index 0000000..40c4f13 --- /dev/null +++ b/tests/integration/test_video_preannotation_integration.py @@ -0,0 +1,596 @@ +""" +Integration tests for video project pre-annotation functionality. + +These tests require real API credentials and will make actual API calls. +Set environment variables: API_KEY, API_SECRET, CLIENT_ID + +Run with: pytest tests/integration/test_video_preannotation_integration.py -m integration +""" + +import json +import os +import tempfile +import time +from concurrent.futures import Future + +import pytest + +from labellerr.core.projects import LabellerrProject + + +@pytest.fixture(scope="module") +def video_annotation_data(): + """Sample video annotation data for testing""" + return [ + { + "file_name": "test_video.mp4", + "annotations": [ + { + "question_name": "Object Detection", + "question_type": "BoundingBox", + "answer": [ + { + "frames": { + "0": { + "frame": 0, + "answer": { + "xmin": 100, + "ymin": 100, + "xmax": 300, + "ymax": 300, + "rotation": 0, + }, + "timestamp": 0.0, + }, + "30": { + "frame": 30, + "answer": { + "xmin": 150, + "ymin": 120, + "xmax": 350, + "ymax": 320, + "rotation": 0, + }, + "timestamp": 1.0, + }, + "60": { + "frame": 60, + "answer": { + "xmin": 180, + "ymin": 150, + "xmax": 380, + "ymax": 350, + "rotation": 0, + }, + "timestamp": 2.0, + }, + }, + }, + { + "frames": { + "0": { + "frame": 0, + "answer": { + "xmin": 500, + "ymin": 300, + "xmax": 800, + "ymax": 450, + "rotation": 0, + }, + "timestamp": 0.0, + }, + "30": { + "frame": 30, + "answer": { + "xmin": 550, + "ymin": 300, + "xmax": 850, + "ymax": 450, + "rotation": 0, + }, + "timestamp": 1.0, + }, + }, + }, + ], + }, + { + "question_name": "Object Tracking", + "question_type": "polygon", + "answer": [ + { + "frames": { + "0": { + "frame": 0, + "answer": [ + {"x": 0, "y": 600}, + {"x": 1920, "y": 600}, + {"x": 1920, "y": 1080}, + {"x": 0, "y": 1080}, + ], + "timestamp": 0.0, + }, + }, + } + ], + }, + ], + } + ] + + +@pytest.fixture(scope="module") +def video_annotation_file(video_annotation_data): + """Create temporary video annotation file""" + temp_file = tempfile.NamedTemporaryFile( + mode="w", suffix=".json", delete=False, encoding="utf-8" + ) + json.dump(video_annotation_data, temp_file) + temp_file.close() + + yield temp_file.name + + # Cleanup + try: + os.unlink(temp_file.name) + except OSError: + pass + + +@pytest.fixture(scope="module") +def coco_annotation_file(): + """Create temporary COCO format annotation file""" + coco_data = { + "images": [ + { + "id": 1, + "file_name": "video_frame_001.jpg", + "width": 1920, + "height": 1080, + } + ], + "annotations": [ + { + "id": 1, + "image_id": 1, + "category_id": 1, + "bbox": [100, 100, 200, 200], + "area": 40000, + "iscrowd": 0, + } + ], + "categories": [{"id": 1, "name": "person", "supercategory": "human"}], + } + + temp_file = tempfile.NamedTemporaryFile( + mode="w", suffix=".json", delete=False, encoding="utf-8" + ) + json.dump(coco_data, temp_file) + temp_file.close() + + yield temp_file.name + + # Cleanup + try: + os.unlink(temp_file.name) + except OSError: + pass + + +@pytest.mark.integration +@pytest.mark.slow +class TestVideoProjectPreannotationIntegration: + """Integration tests for video project pre-annotation""" + + @pytest.fixture(autouse=True) + def setup(self, test_credentials, integration_client): + """Setup for each test""" + self.client = integration_client + self.credentials = test_credentials + + def test_upload_preannotation_sync_video_answer_format(self, video_annotation_file): + """Test synchronous upload of video annotations in video_json format""" + # Get video project instance + project_id = "murial_magnificent_swift_54305" + project = LabellerrProject(self.client, project_id) + + # Verify it's a video project + assert project.data_type == "video", "Test project must be a video project" + + # Upload pre-annotations + result = project.upload_preannotation( + annotation_format="video_json", + annotation_file=video_annotation_file, + ) + + # Verify response structure + assert "response" in result + assert "status" in result["response"] + assert result["response"]["status"] in ["completed", "pending", "processing"] + + # If job is completed, verify metadata + if result["response"]["status"] == "completed": + assert "metadata" in result["response"] + metadata = result["response"]["metadata"] + # Verify metadata contains expected fields + assert isinstance(metadata, dict) + + def test_upload_preannotation_sync_with_confidence_bucket( + self, video_annotation_file + ): + """Test synchronous upload with confidence bucket""" + project_id = "murial_magnificent_swift_54305" + project = LabellerrProject(self.client, project_id) + + # Test each confidence bucket + for conf_bucket in ["low", "medium", "high"]: + result = project.upload_preannotation( + annotation_format="video_json", + annotation_file=video_annotation_file, + conf_bucket=conf_bucket, + ) + + assert "response" in result + assert result["response"]["status"] in [ + "completed", + "pending", + "processing", + ] + + # Add small delay between uploads to avoid rate limiting + time.sleep(2) + + def test_upload_preannotation_async_returns_future(self, video_annotation_file): + """Test asynchronous upload returns Future object""" + project_id = "murial_magnificent_swift_54305" + project = LabellerrProject(self.client, project_id) + + # Upload pre-annotations asynchronously + future = project.upload_preannotation_async( + annotation_format="video_json", + annotation_file=video_annotation_file, + ) + + # Verify Future object is returned + assert isinstance(future, Future), "Should return a Future object" + + # Wait for completion with timeout + result = future.result(timeout=120) + + # Verify result + assert "response" in result + assert "status" in result["response"] + assert result["response"]["status"] in ["completed", "pending", "processing"] + + def test_upload_preannotation_async_with_conf_bucket(self, video_annotation_file): + """Test asynchronous upload with confidence bucket""" + project_id = "murial_magnificent_swift_54305" + project = LabellerrProject(self.client, project_id) + + future = project.upload_preannotation_async( + annotation_format="video_json", + annotation_file=video_annotation_file, + conf_bucket="high", + ) + + # Wait for completion + result = future.result(timeout=120) + + assert "response" in result + assert result["response"]["status"] in ["completed", "pending", "processing"] + + def test_upload_preannotation_coco_format(self, coco_annotation_file): + """Test upload with COCO JSON format""" + project_id = "murial_magnificent_swift_54305" + project = LabellerrProject(self.client, project_id) + + result = project.upload_preannotation( + annotation_format="coco_json", + annotation_file=coco_annotation_file, + ) + + assert "response" in result + assert result["response"]["status"] in ["completed", "pending", "processing"] + + def test_upload_preannotation_invalid_format_fails(self, video_annotation_file): + """Test that invalid annotation format raises error""" + project_id = "murial_magnificent_swift_54305" + project = LabellerrProject(self.client, project_id) + + with pytest.raises(Exception): + project.upload_preannotation( + annotation_format="invalid_format", + annotation_file=video_annotation_file, + ) + + def test_upload_preannotation_invalid_conf_bucket_fails( + self, video_annotation_file + ): + """Test that invalid confidence bucket raises error""" + project_id = "murial_magnificent_swift_54305" + project = LabellerrProject(self.client, project_id) + + with pytest.raises(AssertionError, match="Invalid confidence bucket"): + project._upload_preannotation_sync( + project.project_id, + self.credentials["client_id"], + "video_json", + video_annotation_file, + "invalid_bucket", + ) + + def test_upload_preannotation_nonexistent_file_fails(self): + """Test that nonexistent file raises error""" + project_id = "murial_magnificent_swift_54305" + project = LabellerrProject(self.client, project_id) + + with pytest.raises(Exception): + project.upload_preannotation( + annotation_format="video_json", + annotation_file="/nonexistent/path/to/file.json", + ) + + @pytest.mark.parametrize( + "annotation_format", + ["video_json", "coco_json", "json"], + ) + def test_upload_preannotation_multiple_formats( + self, video_annotation_file, annotation_format + ): + """Test upload with multiple annotation formats""" + project_id = "murial_magnificent_swift_54305" + project = LabellerrProject(self.client, project_id) + + result = project.upload_preannotation( + annotation_format=annotation_format, + annotation_file=video_annotation_file, + ) + + assert "response" in result + assert result["response"]["status"] in ["completed", "pending", "processing"] + + # Add delay between uploads + time.sleep(2) + + +@pytest.mark.integration +@pytest.mark.slow +class TestVideoPreannotationJobStatusMonitoring: + """Integration tests for pre-annotation job status monitoring""" + + @pytest.fixture(autouse=True) + def setup(self, test_credentials, integration_client): + """Setup for each test""" + self.client = integration_client + self.credentials = test_credentials + + def test_job_status_monitoring_completes(self, video_annotation_file): + """Test that job status monitoring waits for completion""" + project_id = "murial_magnificent_swift_54305" + project = LabellerrProject(self.client, project_id) + + # Start async upload + future = project.upload_preannotation_async( + annotation_format="video_json", + annotation_file=video_annotation_file, + ) + + # Monitor status - this should block until complete + start_time = time.time() + result = future.result(timeout=120) + elapsed_time = time.time() - start_time + + # Verify completion + assert result["response"]["status"] in ["completed", "pending", "processing"] + + # If completed, verify it took some time (indicating actual processing) + if result["response"]["status"] == "completed": + # Processing should take at least 1 second + assert elapsed_time >= 1.0, "Job completed too quickly, may not be real" + + def test_concurrent_uploads(self, video_annotation_file): + """Test multiple concurrent pre-annotation uploads""" + project_id = "murial_magnificent_swift_54305" + project = LabellerrProject(self.client, project_id) + + # Start multiple async uploads + futures = [] + for i in range(3): + future = project.upload_preannotation_async( + annotation_format="video_json", + annotation_file=video_annotation_file, + ) + futures.append(future) + time.sleep(1) # Small delay between requests + + # Wait for all to complete + results = [] + for future in futures: + try: + result = future.result(timeout=120) + results.append(result) + except Exception as e: + pytest.fail(f"Concurrent upload failed: {str(e)}") + + # Verify all completed + assert len(results) == 3 + for result in results: + assert "response" in result + assert result["response"]["status"] in [ + "completed", + "pending", + "processing", + ] + + +@pytest.mark.integration +class TestVideoPreannotationErrorHandling: + """Integration tests for error handling in video pre-annotation""" + + @pytest.fixture(autouse=True) + def setup(self, test_credentials, integration_client): + """Setup for each test""" + self.client = integration_client + self.credentials = test_credentials + + def test_upload_with_malformed_json_fails(self): + """Test that malformed JSON file raises appropriate error""" + project_id = "murial_magnificent_swift_54305" + project = LabellerrProject(self.client, project_id) + + # Create malformed JSON file + temp_file = tempfile.NamedTemporaryFile( + mode="w", suffix=".json", delete=False, encoding="utf-8" + ) + temp_file.write("{ invalid json content }") + temp_file.close() + + try: + # This might fail during upload or validation + with pytest.raises(Exception): + project.upload_preannotation( + annotation_format="video_json", + annotation_file=temp_file.name, + ) + finally: + os.unlink(temp_file.name) + + def test_upload_with_empty_file_fails(self): + """Test that empty file raises appropriate error""" + project_id = "murial_magnificent_swift_54305" + project = LabellerrProject(self.client, project_id) + + # Create empty file + temp_file = tempfile.NamedTemporaryFile( + mode="w", suffix=".json", delete=False, encoding="utf-8" + ) + temp_file.close() + + try: + with pytest.raises(Exception): + project.upload_preannotation( + annotation_format="video_json", + annotation_file=temp_file.name, + ) + finally: + os.unlink(temp_file.name) + + def test_upload_to_nonexistent_project_fails(self, video_annotation_file): + """Test upload to non-existent project raises error""" + with pytest.raises(Exception): # type: ignore[misc] + fake_project = LabellerrProject(self.client, "nonexistent_project_id_12345") + fake_project.upload_preannotation( + annotation_format="video_json", + annotation_file=video_annotation_file, + ) + + +@pytest.mark.integration +class TestVideoPreannotationDataFormats: + """Integration tests for various video pre-annotation data formats""" + + @pytest.fixture(autouse=True) + def setup(self, test_credentials, integration_client): + """Setup for each test""" + self.client = integration_client + self.credentials = test_credentials + + def test_upload_single_frame_annotation(self): + """Test upload of single frame annotation""" + project_id = "murial_magnificent_swift_54305" + project = LabellerrProject(self.client, project_id) + + single_frame_data = [ + { + "file_name": "single_frame.mp4", + "annotations": [ + { + "question_name": "Detection", + "question_type": "BoundingBox", + "answer": [ + { + "frames": { + "0": { + "frame": 0, + "answer": { + "xmin": 100, + "ymin": 100, + "xmax": 150, + "ymax": 150, + "rotation": 0, + }, + "timestamp": 0.0, + } + }, + } + ], + } + ], + } + ] + + temp_file = tempfile.NamedTemporaryFile( + mode="w", suffix=".json", delete=False, encoding="utf-8" + ) + json.dump(single_frame_data, temp_file) + temp_file.close() + + try: + result = project.upload_preannotation( + annotation_format="video_json", + annotation_file=temp_file.name, + ) + assert "response" in result + finally: + os.unlink(temp_file.name) + + def test_upload_multi_object_tracking(self): + """Test upload of multiple object tracking annotations""" + project_id = "murial_magnificent_swift_54305" + project = LabellerrProject(self.client, project_id) + + multi_object_data = [ + { + "file_name": "multi_object.mp4", + "annotations": [ + { + "question_name": "Multi-Object Tracking", + "question_type": "BoundingBox", + "answer": [ + { + "frames": { + str(frame): { + "frame": frame, + "answer": { + "xmin": 100 + i * 50, + "ymin": 100 + frame * 10, + "xmax": 150 + i * 50, + "ymax": 150 + frame * 10, + "rotation": 0, + }, + "timestamp": frame / 30.0, + } + for frame in range(0, 60, 10) + }, + } + for i in range(3) + ], + } + ], + } + ] + + temp_file = tempfile.NamedTemporaryFile( + mode="w", suffix=".json", delete=False, encoding="utf-8" + ) + json.dump(multi_object_data, temp_file) + temp_file.close() + + try: + result = project.upload_preannotation( + annotation_format="video_json", + annotation_file=temp_file.name, + ) + assert "response" in result + finally: + os.unlink(temp_file.name) diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 5221dc2..96df217 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -5,8 +5,6 @@ in isolation using mocks and fixtures. """ -import os - import pytest from pydantic import ValidationError @@ -14,6 +12,10 @@ from labellerr.core.projects import create_project from labellerr.core.projects.image_project import ImageProject from labellerr.core.users.base import LabellerrUsers +from labellerr.core.schemas.projects import CreateProjectParams, RotationConfig +from labellerr.core.annotation_templates import LabellerrAnnotationTemplate +from labellerr.core.datasets import LabellerrDataset +from unittest.mock import Mock, patch @pytest.fixture @@ -29,7 +31,7 @@ def project(client): proj = ImageProject.__new__(ImageProject) proj.client = client proj.project_id = "test_project_id" - proj.project_data = project_data + proj.__project_data = project_data return proj @@ -41,151 +43,262 @@ def users(client): @pytest.fixture -def sample_valid_payload(): - """Create a sample valid payload for create_project""" - current_dir = os.path.dirname(os.path.abspath(__file__)) - test_image = os.path.join(current_dir, "test_data", "test_image.jpg") - - # Create test directory and file if they don't exist - os.makedirs(os.path.join(current_dir, "test_data"), exist_ok=True) - if not os.path.exists(test_image): - with open(test_image, "w") as f: - f.write("dummy image content") - - return { - "data_type": "image", - "created_by": "test_user@example.com", - "project_name": "Test Project", - "autolabel": False, - "files_to_upload": [test_image], - "annotation_guide": [ - { - "option_type": "radio", - "question": "Test Question", - "options": ["Option 1", "Option 2"], - } - ], - "rotation_config": { - "annotation_rotation_count": 1, - "review_rotation_count": 1, - "client_review_rotation_count": 1, - }, - } - - -@pytest.mark.unit -class TestInitiateCreateProject: +def sample_valid_params(): + """Create a sample valid CreateProjectParams for create_project""" + return CreateProjectParams( + project_name="Test Project", + data_type="image", + rotations=RotationConfig( + annotation_rotation_count=1, + review_rotation_count=1, + client_review_rotation_count=1, + ), + use_ai=False, + created_by="test_user@example.com", + ) - def test_missing_required_parameters(self, client, sample_valid_payload): - """Test error handling for missing required parameters""" - # Remove required parameters one by one and test - # Current required params in create_project: data_type, created_by, project_name, autolabel - required_params = [ - "data_type", - "created_by", - "project_name", - "autolabel", - ] - for param in required_params: - invalid_payload = sample_valid_payload.copy() - del invalid_payload[param] - - with pytest.raises(LabellerrError) as exc_info: - create_project(client, invalid_payload) +@pytest.fixture +def mock_dataset(): + """Create a mock dataset for testing""" + dataset = Mock(spec=LabellerrDataset) + dataset.dataset_id = "test-dataset-id" + dataset.files_count = 10 + return dataset - assert f"Required parameter {param} is missing" in str(exc_info.value) - # Test annotation_guide separately since it has special validation - invalid_payload = sample_valid_payload.copy() - del invalid_payload["annotation_guide"] +@pytest.fixture +def mock_annotation_template(): + """Create a mock annotation template for testing""" + template = Mock(spec=LabellerrAnnotationTemplate) + template.annotation_template_id = "test-template-id" + return template - with pytest.raises(LabellerrError) as exc_info: - create_project(client, invalid_payload) - assert ( - "Please provide either annotation guide or annotation template id" - in str(exc_info.value) - ) +@pytest.mark.unit +class TestInitiateCreateProject: - def test_invalid_created_by_email(self, client, sample_valid_payload): + def test_missing_required_parameters( + self, client, mock_dataset, mock_annotation_template + ): + """Test error handling for missing required parameters""" + # Test missing project_name + with pytest.raises(ValidationError): + CreateProjectParams( + data_type="image", + rotations=RotationConfig( + annotation_rotation_count=1, + review_rotation_count=1, + client_review_rotation_count=1, + ), + use_ai=False, + created_by="test_user@example.com", + # Missing project_name + ) + + # Test missing data_type + with pytest.raises(ValidationError): + CreateProjectParams( + project_name="Test Project", + rotations=RotationConfig( + annotation_rotation_count=1, + review_rotation_count=1, + client_review_rotation_count=1, + ), + use_ai=False, + created_by="test_user@example.com", + # Missing data_type + ) + + # Test missing rotations + with pytest.raises(ValidationError): + CreateProjectParams( + project_name="Test Project", + data_type="image", + use_ai=False, + created_by="test_user@example.com", + # Missing rotations + ) + + def test_invalid_created_by_email( + self, client, mock_dataset, mock_annotation_template + ): """Test error handling for invalid created_by email format""" - invalid_payload = sample_valid_payload.copy() - invalid_payload["created_by"] = "not_an_email" # Missing @ and domain - - with pytest.raises(LabellerrError) as exc_info: - create_project(client, invalid_payload) - - assert "Please enter email id in created_by" in str(exc_info.value) + # Test invalid email format - should raise ValidationError during CreateProjectParams creation + with pytest.raises(ValidationError): + CreateProjectParams( + project_name="Test Project", + data_type="image", + rotations=RotationConfig( + annotation_rotation_count=1, + review_rotation_count=1, + client_review_rotation_count=1, + ), + use_ai=False, + created_by="not_an_email", # Invalid email format + ) # Test invalid email without domain extension - invalid_payload["created_by"] = "test@example" - with pytest.raises(LabellerrError) as exc_info: - create_project(client, invalid_payload) - - assert "Please enter email id in created_by" in str(exc_info.value) - - def test_invalid_annotation_guide(self, client, sample_valid_payload): + with pytest.raises(ValidationError): + CreateProjectParams( + project_name="Test Project", + data_type="image", + rotations=RotationConfig( + annotation_rotation_count=1, + review_rotation_count=1, + client_review_rotation_count=1, + ), + use_ai=False, + created_by="test@example", # Invalid email format + ) + + def test_invalid_annotation_guide( + self, client, mock_dataset, mock_annotation_template + ): """Test error handling for invalid annotation guide""" - invalid_payload = sample_valid_payload.copy() - - # Missing option_type - invalid_payload["annotation_guide"] = [{"question": "Test Question"}] - with pytest.raises(LabellerrError) as exc_info: - create_project(client, invalid_payload) - - assert "option_type is required in annotation_guide" in str(exc_info.value) - - # Invalid option_type - invalid_payload["annotation_guide"] = [ - {"option_type": "invalid_type", "question": "Test Question"} - ] - with pytest.raises(LabellerrError) as exc_info: - create_project(client, invalid_payload) - - assert "option_type must be one of" in str(exc_info.value) - - def test_both_upload_methods_specified(self, client, sample_valid_payload): - """Test error when both files_to_upload and folder_to_upload are specified""" - invalid_payload = sample_valid_payload.copy() - invalid_payload["folder_to_upload"] = "/path/to/folder" + # Since annotation templates are now separate objects, + # we test that empty datasets raise an error + valid_params = CreateProjectParams( + project_name="Test Project", + data_type="image", + rotations=RotationConfig( + annotation_rotation_count=1, + review_rotation_count=1, + client_review_rotation_count=1, + ), + use_ai=False, + created_by="test_user@example.com", + ) + # Test with empty datasets list with pytest.raises(LabellerrError) as exc_info: - create_project(client, invalid_payload) - - assert "Cannot provide both files_to_upload and folder_to_upload" in str( - exc_info.value + create_project(client, valid_params, [], mock_annotation_template) + + assert "At least one dataset is required" in str(exc_info.value) + + def test_both_upload_methods_specified( + self, client, mock_dataset, mock_annotation_template + ): + """Test error when dataset has no files""" + valid_params = CreateProjectParams( + project_name="Test Project", + data_type="image", + rotations=RotationConfig( + annotation_rotation_count=1, + review_rotation_count=1, + client_review_rotation_count=1, + ), + use_ai=False, + created_by="test_user@example.com", ) - def test_no_upload_method_specified(self, client, sample_valid_payload): - """Test error when neither files_to_upload nor folder_to_upload are specified""" - invalid_payload = sample_valid_payload.copy() - del invalid_payload["files_to_upload"] + # Create a dataset with no files + empty_dataset = Mock(spec=LabellerrDataset) + empty_dataset.dataset_id = "empty-dataset-id" + empty_dataset.files_count = 0 with pytest.raises(LabellerrError) as exc_info: - create_project(client, invalid_payload) - - assert "Either files_to_upload or folder_to_upload must be provided" in str( - exc_info.value + create_project( + client, valid_params, [empty_dataset], mock_annotation_template + ) + + assert "Dataset empty-dataset-id has no files" in str(exc_info.value) + + def test_no_upload_method_specified( + self, client, mock_dataset, mock_annotation_template + ): + """Test successful project creation with valid parameters""" + valid_params = CreateProjectParams( + project_name="Test Project", + data_type="image", + rotations=RotationConfig( + annotation_rotation_count=1, + review_rotation_count=1, + client_review_rotation_count=1, + ), + use_ai=False, + created_by="test_user@example.com", ) - def test_empty_files_to_upload(self, client, sample_valid_payload): - """Test error handling for empty files_to_upload""" - invalid_payload = sample_valid_payload.copy() - invalid_payload["files_to_upload"] = [] - with pytest.raises(LabellerrError): - create_project(client, invalid_payload) - - def test_invalid_folder_to_upload(self, client, sample_valid_payload): - """Test error handling for invalid folder_to_upload""" - invalid_payload = sample_valid_payload.copy() - del invalid_payload["files_to_upload"] - invalid_payload["folder_to_upload"] = " " + # Mock the API response + mock_response = {"response": {"project_id": "test-project-id"}} + + with patch.object(client, "make_request", return_value=mock_response): + with patch( + "labellerr.core.projects.base.LabellerrProject.get_project", + return_value={"project_id": "test-project-id", "data_type": "image"}, + ): + result = create_project( + client, valid_params, [mock_dataset], mock_annotation_template + ) + assert result is not None + + def test_empty_files_to_upload( + self, client, mock_dataset, mock_annotation_template + ): + """Test project creation with multiple datasets""" + valid_params = CreateProjectParams( + project_name="Test Project", + data_type="image", + rotations=RotationConfig( + annotation_rotation_count=1, + review_rotation_count=1, + client_review_rotation_count=1, + ), + use_ai=False, + created_by="test_user@example.com", + ) - with pytest.raises(LabellerrError) as exc_info: - create_project(client, invalid_payload) + # Create multiple datasets + dataset1 = Mock(spec=LabellerrDataset) + dataset1.dataset_id = "dataset-1" + dataset1.files_count = 5 + + dataset2 = Mock(spec=LabellerrDataset) + dataset2.dataset_id = "dataset-2" + dataset2.files_count = 10 + + # Mock the API response + mock_response = {"response": {"project_id": "test-project-id"}} + + with patch.object(client, "make_request", return_value=mock_response): + with patch( + "labellerr.core.projects.base.LabellerrProject.get_project", + return_value={"project_id": "test-project-id", "data_type": "image"}, + ): + result = create_project( + client, valid_params, [dataset1, dataset2], mock_annotation_template + ) + assert result is not None + + def test_invalid_folder_to_upload( + self, client, mock_dataset, mock_annotation_template + ): + """Test project creation with AI enabled""" + valid_params = CreateProjectParams( + project_name="Test Project", + data_type="image", + rotations=RotationConfig( + annotation_rotation_count=2, + review_rotation_count=2, + client_review_rotation_count=1, + ), + use_ai=True, # Enable AI + created_by="test_user@example.com", + ) - assert "Folder path does not exist" in str(exc_info.value) + # Mock the API response + mock_response = {"response": {"project_id": "test-project-id"}} + + with patch.object(client, "make_request", return_value=mock_response): + with patch( + "labellerr.core.projects.base.LabellerrProject.get_project", + return_value={"project_id": "test-project-id", "data_type": "image"}, + ): + result = create_project( + client, valid_params, [mock_dataset], mock_annotation_template + ) + assert result is not None @pytest.mark.unit diff --git a/tests/unit/test_create_dataset_path.py b/tests/unit/test_create_dataset_path.py index 6bfa3f6..eee98a7 100644 --- a/tests/unit/test_create_dataset_path.py +++ b/tests/unit/test_create_dataset_path.py @@ -1,93 +1,32 @@ """ -Unit tests for create_dataset path parameter validation. +Unit tests for dataset creation functions. -This module focuses on testing path parameter handling for AWS and GCS connectors -in the create_dataset functionality. +This module focuses on testing the create_dataset_from_connection and +create_dataset_from_local functionality. """ from unittest.mock import patch import pytest -from labellerr.core.datasets import create_dataset +from labellerr.core.datasets import ( + create_dataset_from_connection, + create_dataset_from_local, +) +from labellerr.core.datasets.base import LabellerrDataset from labellerr.core.exceptions import LabellerrError from labellerr.core.schemas import DatasetConfig @pytest.mark.unit -class TestCreateDatasetPathValidation: - """Test path parameter validation for AWS and GCS connectors""" +class TestCreateDatasetFunctions: + """Test dataset creation functions""" - def test_aws_connector_missing_path_with_config(self, client): - """Test that AWS connector requires path when using connector_config""" + def test_create_dataset_from_connection_success(self, client): + """Test successful dataset creation from connection""" dataset_config = DatasetConfig( client_id="test_client_id", - dataset_name="Test AWS Dataset", - data_type="image", - connector_type="aws", - ) - - aws_config = { - "aws_access_key_id": "test-key", - "aws_secret_access_key": "test-secret", - "aws_region": "us-east-1", - "bucket_name": "test-bucket", - "data_type": "image", - } - - with pytest.raises(LabellerrError) as exc_info: - create_dataset( - client=client, - dataset_config=dataset_config, - connector_config=aws_config, - # Missing path parameter - ) - - assert "path is required for aws connector" in str(exc_info.value) - - def test_gcp_connector_missing_path_with_config(self, client): - """Test that GCP connector requires path when using connector_config""" - dataset_config = DatasetConfig( - client_id="test_client_id", - dataset_name="Test GCP Dataset", - data_type="image", - connector_type="gcp", - ) - - # Create a temporary credentials file for testing - import json - import tempfile - - with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f: - json.dump({"type": "service_account"}, f) - temp_cred_file = f.name - - try: - gcp_config = { - "gcs_cred_file": temp_cred_file, - "gcs_path": "gs://test-bucket/path", - "data_type": "image", - } - - with pytest.raises(LabellerrError) as exc_info: - create_dataset( - client=client, - dataset_config=dataset_config, - connector_config=gcp_config, - # Missing path parameter - ) - - assert "path is required for gcp connector" in str(exc_info.value) - finally: - import os - - os.unlink(temp_cred_file) - - def test_aws_connector_with_path_and_connection_id(self, client): - """Test AWS connector with both path and existing connection_id""" - dataset_config = DatasetConfig( - client_id="test_client_id", - dataset_name="Test AWS Dataset", + dataset_name="Test Dataset", data_type="image", connector_type="aws", ) @@ -102,47 +41,65 @@ def test_aws_connector_with_path_and_connection_id(self, client): "labellerr.core.datasets.base.LabellerrDataset.get_dataset", return_value={"dataset_id": "test-dataset-id", "data_type": "image"}, ): - dataset = create_dataset( + # Create a mock connection object + from unittest.mock import Mock + + mock_connection = Mock() + mock_connection.connection_id = "test-connection-id" + + dataset = create_dataset_from_connection( client=client, dataset_config=dataset_config, + connection=mock_connection, path="s3://test-bucket/path/to/data", - connection_id="existing-aws-connection-id", ) - # Should succeed - path is provided with connection_id + # Should succeed assert dataset is not None - def test_gcp_connector_with_path_and_connection_id(self, client): - """Test GCP connector with both path and existing connection_id""" + def test_create_dataset_from_local_with_files(self, client): + """Test successful dataset creation from local files""" dataset_config = DatasetConfig( client_id="test_client_id", - dataset_name="Test GCP Dataset", + dataset_name="Test Local Dataset", data_type="image", - connector_type="gcp", + connector_type="local", ) mock_response = { "response": {"dataset_id": "test-dataset-id", "data_type": "image"} } - # Mock both the dataset creation and the get_dataset call - with patch.object(client, "make_request", return_value=mock_response): - with patch( - "labellerr.core.datasets.base.LabellerrDataset.get_dataset", - return_value={"dataset_id": "test-dataset-id", "data_type": "image"}, - ): - dataset = create_dataset( - client=client, - dataset_config=dataset_config, - path="gs://test-bucket/path/to/data", - connection_id="existing-gcp-connection-id", - ) - - # Should succeed - path is provided with connection_id - assert dataset is not None - - def test_local_connector_path_parameter_ignored(self, client): - """Test that local connector ignores path parameter""" + # Mock the upload_files function and dataset creation + with patch( + "labellerr.core.datasets.upload_files", return_value="test-connection-id" + ): + with patch.object(client, "make_request", return_value=mock_response): + with patch( + "labellerr.core.datasets.base.LabellerrDataset.get_dataset", + return_value={ + "dataset_id": "test-dataset-id", + "data_type": "image", + }, + ): + # Mock create_dataset_from_connection to avoid the connection object issue + with patch( + "labellerr.core.datasets.create_dataset_from_connection", + return_value=LabellerrDataset( + client=client, dataset_id="test-dataset-id" + ), + ): + dataset = create_dataset_from_local( + client=client, + dataset_config=dataset_config, + files_to_upload=["test_file1.jpg", "test_file2.jpg"], + ) + + # Should succeed + assert dataset is not None + + def test_create_dataset_from_local_with_folder(self, client): + """Test successful dataset creation from local folder""" dataset_config = DatasetConfig( client_id="test_client_id", dataset_name="Test Local Dataset", @@ -154,85 +111,84 @@ def test_local_connector_path_parameter_ignored(self, client): "response": {"dataset_id": "test-dataset-id", "data_type": "image"} } - # Mock both the dataset creation and the get_dataset call - with patch.object(client, "make_request", return_value=mock_response): - with patch( - "labellerr.core.datasets.base.LabellerrDataset.get_dataset", - return_value={"dataset_id": "test-dataset-id", "data_type": "image"}, - ): - dataset = create_dataset( - client=client, - dataset_config=dataset_config, - path="some/ignored/path", # Should be ignored for local - ) - - # Should succeed - path is not validated for local connector - assert dataset is not None - - def test_aws_missing_connector_config_and_connection_id(self, client): - """Test error when neither connector_config nor connection_id is provided for AWS""" + # Mock the upload_folder_files_to_dataset function and dataset creation + with patch( + "labellerr.core.datasets.upload_folder_files_to_dataset", + return_value={"connection_id": "test-connection-id", "status": "success"}, + ): + with patch.object(client, "make_request", return_value=mock_response): + with patch( + "labellerr.core.datasets.base.LabellerrDataset.get_dataset", + return_value={ + "dataset_id": "test-dataset-id", + "data_type": "image", + }, + ): + # Mock create_dataset_from_connection to avoid the connection object issue + with patch( + "labellerr.core.datasets.create_dataset_from_connection", + return_value=LabellerrDataset( + client=client, dataset_id="test-dataset-id" + ), + ): + dataset = create_dataset_from_local( + client=client, + dataset_config=dataset_config, + folder_to_upload="/path/to/test/folder", + ) + + # Should succeed + assert dataset is not None + + def test_create_dataset_from_local_no_files_or_folder(self, client): + """Test error when neither files nor folder is provided for local dataset""" dataset_config = DatasetConfig( client_id="test_client_id", - dataset_name="Test AWS Dataset", + dataset_name="Test Local Dataset", data_type="image", - connector_type="aws", + connector_type="local", ) with pytest.raises(LabellerrError) as exc_info: - create_dataset( + create_dataset_from_local( client=client, dataset_config=dataset_config, - path="s3://test-bucket/path/to/data", - # Missing both connector_config and connection_id + # Missing both files_to_upload and folder_to_upload ) - assert "connector_config is required for aws connector" in str(exc_info.value) + assert "No files or folder to upload provided" in str(exc_info.value) - def test_gcp_missing_connector_config_and_connection_id(self, client): - """Test error when neither connector_config nor connection_id is provided for GCP""" + def test_create_dataset_from_connection_with_different_paths(self, client): + """Test dataset creation with different path formats""" dataset_config = DatasetConfig( client_id="test_client_id", - dataset_name="Test GCP Dataset", + dataset_name="Test Dataset", data_type="image", connector_type="gcp", ) - with pytest.raises(LabellerrError) as exc_info: - create_dataset( - client=client, - dataset_config=dataset_config, - path="gs://test-bucket/path/to/data", - # Missing both connector_config and connection_id - ) - - assert "connector_config is required for gcp connector" in str(exc_info.value) + mock_response = { + "response": {"dataset_id": "test-dataset-id", "data_type": "image"} + } - def test_both_connection_id_and_connector_config(self, client): - """Test error when both connection_id and connector_config are provided""" - dataset_config = DatasetConfig( - client_id="test_client_id", - dataset_name="Test AWS Dataset", - data_type="image", - connector_type="aws", - ) + # Test with GCS path + with patch.object(client, "make_request", return_value=mock_response): + with patch( + "labellerr.core.datasets.base.LabellerrDataset.get_dataset", + return_value={"dataset_id": "test-dataset-id", "data_type": "image"}, + ): + # Create a mock connection object + from unittest.mock import Mock - aws_config = { - "aws_access_key_id": "test-key", - "aws_secret_access_key": "test-secret", - "aws_region": "us-east-1", - "bucket_name": "test-bucket", - "data_type": "image", - } + mock_connection = Mock() + mock_connection.connection_id = "test-gcp-connection-id" - with pytest.raises(LabellerrError) as exc_info: - create_dataset( - client=client, - dataset_config=dataset_config, - path="s3://test-bucket/path/to/data", - connection_id="existing-connection-id", - connector_config=aws_config, - ) + dataset = create_dataset_from_connection( + client=client, + dataset_config=dataset_config, + connection=mock_connection, + path="gs://test-bucket/path/to/data", + ) - assert "Cannot provide both connection_id and connector_config" in str( - exc_info.value - ) + # Should succeed + assert dataset is not None diff --git a/tests/unit/test_dataset_pagination.py b/tests/unit/test_dataset_pagination.py index 5a0abe1..0636327 100644 --- a/tests/unit/test_dataset_pagination.py +++ b/tests/unit/test_dataset_pagination.py @@ -9,7 +9,7 @@ import pytest -from labellerr.core.datasets.base import LabellerrDataset +from labellerr.core.datasets import list_datasets from labellerr.core.schemas import DataSetScope # Helper to use correct enum values @@ -87,9 +87,7 @@ def test_default_page_size_used(self, client, mock_single_page_response): with patch.object(client, "make_request") as mock_request: mock_request.return_value = mock_single_page_response - result = LabellerrDataset.get_all_datasets( - client=client, datatype="image", scope=SCOPE_CLIENT - ) + result = list_datasets(client=client, datatype="image", scope=SCOPE_CLIENT) # Consume the generator to trigger the API call list(result) @@ -109,9 +107,7 @@ def test_default_returns_generator(self, client, mock_single_page_response): with patch.object(client, "make_request") as mock_request: mock_request.return_value = mock_single_page_response - result = LabellerrDataset.get_all_datasets( - client=client, datatype="image", scope=SCOPE_CLIENT - ) + result = list_datasets(client=client, datatype="image", scope=SCOPE_CLIENT) # Check that result is a generator import types @@ -133,7 +129,7 @@ def test_custom_page_size(self, client, mock_single_page_response): with patch.object(client, "make_request") as mock_request: mock_request.return_value = mock_single_page_response - result = LabellerrDataset.get_all_datasets( + result = list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, page_size=20 ) @@ -149,7 +145,7 @@ def test_pagination_with_last_dataset_id(self, client, mock_second_page_response with patch.object(client, "make_request") as mock_request: mock_request.return_value = mock_second_page_response - result = LabellerrDataset.get_all_datasets( + result = list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, @@ -172,7 +168,7 @@ def test_manual_pagination_flow( with patch.object(client, "make_request") as mock_request: # First page mock_request.return_value = mock_first_page_response - first_page_gen = LabellerrDataset.get_all_datasets( + first_page_gen = list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, page_size=2 ) first_page_datasets = list(first_page_gen) @@ -187,7 +183,7 @@ def test_manual_pagination_flow( # Second page mock_request.return_value = mock_last_page_response - second_page_gen = LabellerrDataset.get_all_datasets( + second_page_gen = list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, @@ -209,7 +205,7 @@ def test_auto_pagination_returns_generator(self, client, mock_single_page_respon with patch.object(client, "make_request") as mock_request: mock_request.return_value = mock_single_page_response - result = LabellerrDataset.get_all_datasets( + result = list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, @@ -229,7 +225,7 @@ def test_auto_pagination_yields_individual_datasets( mock_request.return_value = mock_single_page_response datasets = list( - LabellerrDataset.get_all_datasets( + list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, @@ -249,7 +245,7 @@ def test_auto_pagination_single_page(self, client, mock_single_page_response): mock_request.return_value = mock_single_page_response datasets = list( - LabellerrDataset.get_all_datasets( + list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, @@ -278,7 +274,7 @@ def test_auto_pagination_multiple_pages( ] datasets = list( - LabellerrDataset.get_all_datasets( + list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, @@ -305,7 +301,7 @@ def test_auto_pagination_uses_default_page_size_internally( mock_request.return_value = mock_single_page_response list( - LabellerrDataset.get_all_datasets( + list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, @@ -329,7 +325,7 @@ def test_auto_pagination_passes_last_dataset_id( ] list( - LabellerrDataset.get_all_datasets( + list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, @@ -351,7 +347,7 @@ def test_auto_pagination_early_termination( mock_second_page_response, ] - generator = LabellerrDataset.get_all_datasets( + generator = list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, @@ -385,9 +381,7 @@ def test_empty_results(self, client): with patch.object(client, "make_request") as mock_request: mock_request.return_value = empty_response - result = LabellerrDataset.get_all_datasets( - client=client, datatype="image", scope=SCOPE_CLIENT - ) + result = list_datasets(client=client, datatype="image", scope=SCOPE_CLIENT) datasets = list(result) assert datasets == [] @@ -402,7 +396,7 @@ def test_empty_results_auto_pagination(self, client): mock_request.return_value = empty_response datasets = list( - LabellerrDataset.get_all_datasets( + list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, @@ -420,7 +414,7 @@ def test_different_data_types(self, client, mock_single_page_response): with patch.object(client, "make_request") as mock_request: mock_request.return_value = mock_single_page_response - result = LabellerrDataset.get_all_datasets( + result = list_datasets( client=client, datatype=data_type, scope=SCOPE_CLIENT ) list(result) # Consume generator @@ -437,9 +431,7 @@ def test_different_scopes(self, client, mock_single_page_response): with patch.object(client, "make_request") as mock_request: mock_request.return_value = mock_single_page_response - result = LabellerrDataset.get_all_datasets( - client=client, datatype="image", scope=scope - ) + result = list_datasets(client=client, datatype="image", scope=scope) list(result) # Consume generator call_args = mock_request.call_args @@ -451,7 +443,7 @@ def test_large_page_size(self, client, mock_single_page_response): with patch.object(client, "make_request") as mock_request: mock_request.return_value = mock_single_page_response - result = LabellerrDataset.get_all_datasets( + result = list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, page_size=1000 ) list(result) # Consume generator @@ -470,7 +462,7 @@ def test_auto_pagination_memory_efficiency( mock_last_page_response, ] - generator = LabellerrDataset.get_all_datasets( + generator = list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, @@ -500,7 +492,7 @@ def test_iterate_with_for_loop( ] dataset_ids = [] - for dataset in LabellerrDataset.get_all_datasets( + for dataset in list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, @@ -522,7 +514,7 @@ def test_list_comprehension( dataset_names = [ d["name"] - for d in LabellerrDataset.get_all_datasets( + for d in list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, @@ -545,7 +537,7 @@ def test_filtering_while_iterating( # Get only datasets with even IDs even_datasets = [ d - for d in LabellerrDataset.get_all_datasets( + for d in list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, diff --git a/tests/unit/test_keyframes.py b/tests/unit/test_keyframes.py index 132fc42..9650254 100644 --- a/tests/unit/test_keyframes.py +++ b/tests/unit/test_keyframes.py @@ -201,7 +201,7 @@ def mock_video_project(mock_client): project = VideoProject.__new__(VideoProject) project.client = mock_client project.project_id = "test_project_id" - project.project_data = { + project.__project_data = { "project_id": "test_project_id", "data_type": "video", "attached_datasets": [], diff --git a/tests/unit/test_video_preannotation.py b/tests/unit/test_video_preannotation.py new file mode 100644 index 0000000..efe9853 --- /dev/null +++ b/tests/unit/test_video_preannotation.py @@ -0,0 +1,431 @@ +""" +Unit tests for video project pre-annotation functionality. + +This module contains comprehensive unit tests for video project pre-annotation +upload methods, including synchronous and asynchronous operations. +""" + +import json +import os +import tempfile +from concurrent.futures import Future +from unittest.mock import Mock, patch + +import pytest + +from labellerr.client import LabellerrClient +from labellerr.core.exceptions import LabellerrError + + +@pytest.fixture +def mock_client(): + """Create a mock client for testing""" + client = LabellerrClient("test_api_key", "test_api_secret", "test_client_id") + client.base_url = "https://api.labellerr.com" + return client + + +@pytest.fixture +def mock_video_project(mock_client): + """Create a mock video project instance""" + from labellerr.core.projects.video_project import VideoProject + + # Create instance bypassing metaclass + project = VideoProject.__new__(VideoProject) + project.client = mock_client + project.project_id = "test_project_id" + project.project_data = { + "project_id": "test_project_id", + "data_type": "video", + "attached_datasets": [], + } + return project + + +@pytest.fixture +def temp_annotation_file(): + """Create temporary annotation file for testing""" + annotation_data = [ + { + "file_name": "video1.mp4", + "annotations": [ + { + "question_name": "Object Detection", + "question_type": "BoundingBox", + "answer": [ + { + "frames": { + "0": { + "frame": 0, + "answer": { + "xmin": 100, + "ymin": 100, + "xmax": 300, + "ymax": 300, + "rotation": 0, + }, + "timestamp": 0.0, + }, + "25": { + "frame": 25, + "answer": { + "xmin": 150, + "ymin": 150, + "xmax": 350, + "ymax": 350, + "rotation": 0, + }, + "timestamp": 1.0, + }, + }, + } + ], + } + ], + } + ] + + temp_file = tempfile.NamedTemporaryFile( + mode="w", suffix=".json", delete=False, encoding="utf-8" + ) + json.dump(annotation_data, temp_file) + temp_file.close() + + yield temp_file.name + + # Cleanup + try: + os.unlink(temp_file.name) + except OSError: + pass + + +@pytest.mark.unit +class TestVideoProjectPreannotationUpload: + """Unit tests for video project pre-annotation upload methods""" + + def test_upload_preannotation_delegates_to_base_method( + self, mock_video_project, temp_annotation_file + ): + """Test that upload_preannotation delegates to upload_preannotations""" + with patch.object( + mock_video_project, + "upload_preannotations", + return_value={"status": "completed"}, + ) as mock_upload: + result = mock_video_project.upload_preannotation( + "video_json", temp_annotation_file, "medium" + ) + + mock_upload.assert_called_once_with( + "video_json", + temp_annotation_file, + "medium", + _async=False, + ) + assert result == {"status": "completed"} + + def test_upload_preannotation_without_conf_bucket( + self, mock_video_project, temp_annotation_file + ): + """Test upload_preannotation without confidence bucket""" + with patch.object( + mock_video_project, + "upload_preannotations", + return_value={"status": "completed"}, + ) as mock_upload: + result = mock_video_project.upload_preannotation( + "video_json", temp_annotation_file + ) + + mock_upload.assert_called_once_with( + "video_json", + temp_annotation_file, + None, + _async=False, + ) + assert result == {"status": "completed"} + + @pytest.mark.parametrize( + "annotation_format", + ["video_json", "coco_json", "yolo", "json"], + ) + def test_upload_preannotation_various_formats( + self, mock_video_project, temp_annotation_file, annotation_format + ): + """Test upload_preannotation with various annotation formats""" + with patch.object( + mock_video_project, + "upload_preannotations", + return_value={"status": "completed"}, + ) as mock_upload: + mock_video_project.upload_preannotation( + annotation_format, temp_annotation_file + ) + + mock_upload.assert_called_once() + assert mock_upload.call_args[0][0] == annotation_format + + @pytest.mark.parametrize( + "conf_bucket", + ["low", "medium", "high"], + ) + def test_upload_preannotation_various_conf_buckets( + self, mock_video_project, temp_annotation_file, conf_bucket + ): + """Test upload_preannotation with various confidence buckets""" + with patch.object( + mock_video_project, + "upload_preannotations", + return_value={"status": "completed"}, + ) as mock_upload: + mock_video_project.upload_preannotation( + "video_json", temp_annotation_file, conf_bucket + ) + + mock_upload.assert_called_once() + assert mock_upload.call_args[0][2] == conf_bucket + + +@pytest.mark.unit +class TestVideoProjectPreannotationUploadAsync: + """Unit tests for video project async pre-annotation upload""" + + def test_upload_preannotation_async_returns_future( + self, mock_video_project, temp_annotation_file + ): + """Test that upload_preannotation_async returns a Future object""" + mock_future = Mock(spec=Future) + mock_future.result.return_value = {"status": "completed"} + + with patch.object( + mock_video_project, + "upload_preannotations", + return_value=mock_future, + ): + result = mock_video_project.upload_preannotation_async( + "video_json", temp_annotation_file + ) + + assert isinstance(result, Mock) # Mock of Future + + def test_upload_preannotation_async_delegates_to_base( + self, mock_video_project, temp_annotation_file + ): + """Test that async method delegates to upload_preannotations with _async=True""" + mock_future = Mock(spec=Future) + + with patch.object( + mock_video_project, + "upload_preannotations", + return_value=mock_future, + ) as mock_base_async: + result = mock_video_project.upload_preannotation_async( + "video_json", temp_annotation_file, "high" + ) + + mock_base_async.assert_called_once_with( + "video_json", temp_annotation_file, "high", _async=True + ) + assert result == mock_future + + def test_upload_preannotation_async_without_conf_bucket( + self, mock_video_project, temp_annotation_file + ): + """Test async upload without confidence bucket""" + mock_future = Mock(spec=Future) + + with patch.object( + mock_video_project, + "upload_preannotations", + return_value=mock_future, + ) as mock_base_async: + mock_video_project.upload_preannotation_async( + "video_json", temp_annotation_file + ) + + mock_base_async.assert_called_once_with( + "video_json", temp_annotation_file, None, _async=True + ) + + @pytest.mark.parametrize( + "annotation_format,conf_bucket", + [ + ("video_json", "low"), + ("coco_json", "medium"), + ("json", "high"), + ("yolo", None), + ], + ) + def test_upload_preannotation_async_various_combinations( + self, mock_video_project, temp_annotation_file, annotation_format, conf_bucket + ): + """Test async upload with various format and confidence bucket combinations""" + mock_future = Mock(spec=Future) + + with patch.object( + mock_video_project, + "upload_preannotations", + return_value=mock_future, + ) as mock_base_async: + mock_video_project.upload_preannotation_async( + annotation_format, temp_annotation_file, conf_bucket + ) + + mock_base_async.assert_called_once_with( + annotation_format, temp_annotation_file, conf_bucket, _async=True + ) + + +@pytest.mark.unit +class TestPreannotationSyncInternal: + """Unit tests for preannotation upload internal behavior""" + + def test_upload_preannotation_with_valid_file( + self, + mock_video_project, + temp_annotation_file, + ): + """Test upload with valid annotation file""" + with patch.object( + mock_video_project, + "upload_preannotations", + return_value={"status": "completed", "response": {"job_id": "job-123"}}, + ) as mock_upload: + result = mock_video_project.upload_preannotation( + "video_json", + temp_annotation_file, + "medium", + ) + + mock_upload.assert_called_once() + assert result["status"] == "completed" + + def test_upload_preannotation_sync_missing_params( + self, mock_video_project, temp_annotation_file + ): + """Test sync upload with missing required parameters""" + with patch.object( + mock_video_project, + "upload_preannotations", + side_effect=LabellerrError("Missing required parameter"), + ): + with pytest.raises(LabellerrError): + mock_video_project.upload_preannotation( + None, # Missing annotation_format + temp_annotation_file, + ) + + def test_upload_preannotation_sync_with_conf_bucket_url( + self, + mock_video_project, + temp_annotation_file, + ): + """Test that conf_bucket is properly passed""" + with patch.object( + mock_video_project, + "upload_preannotations", + return_value={"status": "completed", "response": {"job_id": "job-123"}}, + ) as mock_upload: + mock_video_project.upload_preannotation( + "video_json", + temp_annotation_file, + "high", + ) + + # Verify conf_bucket was passed correctly + call_args = mock_upload.call_args + assert call_args[0][2] == "high" + + def test_upload_preannotation_sync_invalid_conf_bucket( + self, + mock_video_project, + temp_annotation_file, + ): + """Test sync upload with invalid confidence bucket""" + with patch.object( + mock_video_project, + "upload_preannotations", + side_effect=AssertionError("Invalid confidence bucket"), + ): + with pytest.raises(AssertionError): + mock_video_project.upload_preannotation( + "video_json", + temp_annotation_file, + "invalid_bucket", # Invalid bucket + ) + + +@pytest.mark.unit +class TestPreannotationFileValidation: + """Unit tests for pre-annotation file validation""" + + def test_video_json_format_with_json_file(self, mock_video_project): + """Test that video_json format accepts JSON files""" + temp_file = tempfile.NamedTemporaryFile( + mode="w", suffix=".json", delete=False, encoding="utf-8" + ) + temp_file.write('{"test": "data"}') + temp_file.close() + + try: + with patch.object( + mock_video_project, + "upload_preannotations", + return_value={"status": "completed"}, + ): + # Should not raise any error + mock_video_project.upload_preannotation("video_json", temp_file.name) + finally: + os.unlink(temp_file.name) + + def test_nonexistent_file(self, mock_video_project): + """Test upload with nonexistent file""" + with patch.object( + mock_video_project, + "upload_preannotations", + side_effect=LabellerrError("File not found"), + ): + with pytest.raises(LabellerrError): + mock_video_project.upload_preannotation( + "video_json", + "/nonexistent/file.json", + ) + + +@pytest.mark.unit +class TestPreannotationJobStatus: + """Unit tests for pre-annotation job status monitoring""" + + @patch("labellerr.core.projects.base.concurrent.futures.ThreadPoolExecutor") + def test_preannotation_job_status_async_calls_poll( + self, mock_executor, mock_video_project + ): + """Test that job status monitoring uses polling mechanism""" + mock_future = Mock(spec=Future) + mock_executor.return_value.__enter__.return_value.submit.return_value = ( + mock_future + ) + + result = mock_video_project.preannotation_job_status_async("job-123") + + # Verify executor was used and Future was returned + assert result == mock_future + + @patch("labellerr.core.projects.base.concurrent.futures.ThreadPoolExecutor") + def test_async_upload_uses_thread_pool( + self, mock_executor, mock_video_project, temp_annotation_file + ): + """Test that async upload returns a Future""" + mock_future = Mock(spec=Future) + + with patch.object( + mock_video_project, + "upload_preannotations", + return_value=mock_future, + ): + result = mock_video_project.upload_preannotation_async( + "video_json", temp_annotation_file + ) + # Verify it returns the future + assert result == mock_future