From 7b5b3940e45223a6e40e17c649eaf2002aeec3c5 Mon Sep 17 00:00:00 2001 From: Gaurav Dudeja Date: Mon, 5 Jan 2026 23:24:52 +0530 Subject: [PATCH] sam2 trigger via sdk --- labellerr/core/client.py | 2 ++ labellerr/core/sam2/__init__.py | 3 ++ labellerr/core/sam2/base.py | 42 ++++++++++++++++++++++ labellerr/core/schemas/datasets.py | 4 +-- tests/integration/test_sam2_integration.py | 34 ++++++++++++++++++ tests/unit/test_sam2.py | 40 +++++++++++++++++++++ 6 files changed, 123 insertions(+), 2 deletions(-) create mode 100644 labellerr/core/sam2/__init__.py create mode 100644 labellerr/core/sam2/base.py create mode 100644 tests/integration/test_sam2_integration.py create mode 100644 tests/unit/test_sam2.py diff --git a/labellerr/core/client.py b/labellerr/core/client.py index 8659bc1..8cdf03d 100644 --- a/labellerr/core/client.py +++ b/labellerr/core/client.py @@ -50,9 +50,11 @@ def __init__( # Import here to avoid circular imports from .users.base import LabellerrUsers + from .sam2.base import LabellerrSam2 # Initialize Users handler for user-related operations self.users = LabellerrUsers(self) + self.sam2 = LabellerrSam2(self) def _setup_session(self): """ diff --git a/labellerr/core/sam2/__init__.py b/labellerr/core/sam2/__init__.py new file mode 100644 index 0000000..f0f7a62 --- /dev/null +++ b/labellerr/core/sam2/__init__.py @@ -0,0 +1,3 @@ +from .base import LabellerrSam2 + +__all__ = ["LabellerrSam2"] diff --git a/labellerr/core/sam2/base.py b/labellerr/core/sam2/base.py new file mode 100644 index 0000000..a9fb735 --- /dev/null +++ b/labellerr/core/sam2/base.py @@ -0,0 +1,42 @@ +import uuid +from typing import TYPE_CHECKING, Dict, Any + +if TYPE_CHECKING: + from ..client import LabellerrClient + + +class LabellerrSam2: + def __init__(self, client: "LabellerrClient"): + self.client = client + + def create_job_from_annotations( + self, project_id: str, file_id: str, email_id: str + ) -> Dict[str, Any]: + """ + Create a SAM2 job from annotations. + + :param project_id: The ID of the project. + :param file_id: The ID of the file. + :param email_id: The email ID of the user. + :return: A dictionary containing job_ids and a message. + """ + unique_id = str(uuid.uuid4()) + url = ( + f"{self.client.base_url}/sam2/create_job_from_annotations" + f"?client_id={self.client.client_id}&project_id={project_id}&uuid={unique_id}" + ) + + payload = { + "file_id": file_id, + "project_id": project_id, + "email_id": email_id, + } + + response = self.client.make_request( + "POST", + url, + extra_headers={"content-type": "application/json"}, + request_id=unique_id, + json=payload, + ) + return response.get("response", None) diff --git a/labellerr/core/schemas/datasets.py b/labellerr/core/schemas/datasets.py index 2ada187..6127beb 100644 --- a/labellerr/core/schemas/datasets.py +++ b/labellerr/core/schemas/datasets.py @@ -3,7 +3,7 @@ """ import os -from enum import StrEnum +from enum import Enum from typing import List, Literal from uuid import UUID @@ -11,7 +11,7 @@ from pydantic import BaseModel, Field, field_validator -class DataSetScope(StrEnum): +class DataSetScope(str, Enum): project = "project" client = "client" public = "public" diff --git a/tests/integration/test_sam2_integration.py b/tests/integration/test_sam2_integration.py new file mode 100644 index 0000000..44e86bb --- /dev/null +++ b/tests/integration/test_sam2_integration.py @@ -0,0 +1,34 @@ +import pytest +import os +from labellerr.core.exceptions import LabellerrError + +@pytest.mark.integration +class TestSam2Workflow: + """Test SAM2 workflow""" + + def test_create_job_from_annotations(self, integration_client): + """Test creating a SAM2 job from annotations""" + + # Use IDs provided by user or from env, with fallbacks + project_id = os.getenv("TEST_SAM2_PROJECT_ID", "ninnetta_necessary_penguin_93195") + file_id = os.getenv("TEST_SAM2_FILE_ID", "2a8d96ca-9161-4dee-ad3b-a5faf301bc6c") + email_id = os.getenv("TEST_SAM2_EMAIL_ID", "dev@labellerr.com") + + try: + response = integration_client.sam2.create_job_from_annotations( + project_id=project_id, + file_id=file_id, + email_id=email_id + ) + + assert isinstance(response, dict) + assert "job_ids" in response + assert "message" in response + assert isinstance(response["job_ids"], list) + + except LabellerrError as e: + # Handle cases where the specific project/file might not exist in the test env + if any(phrase in str(e).lower() for phrase in ["not found", "permission denied", "404", "403"]): + pytest.skip(f"Skipping SAM2 test due to invalid resource or permission: {e}") + else: + raise diff --git a/tests/unit/test_sam2.py b/tests/unit/test_sam2.py new file mode 100644 index 0000000..0094e18 --- /dev/null +++ b/tests/unit/test_sam2.py @@ -0,0 +1,40 @@ +import pytest +from unittest.mock import Mock, patch +from labellerr.core.sam2.base import LabellerrSam2 +from labellerr.client import LabellerrClient + +class TestSam2Unit: + def test_create_job_from_annotations_payload(self): + """Test that create_job_from_annotations sends correct payload""" + mock_client = Mock(spec=LabellerrClient) + mock_client.base_url = "https://api.labellerr.com" + mock_client.client_id = "test_client_id" + mock_client.make_request.return_value = { + "response": { + "job_ids": ["job_123"], + "message": "Created 1 jobs" + } + } + + sam2 = LabellerrSam2(mock_client) + + project_id = "test_project" + file_id = "test_file" + email_id = "test@example.com" + + result = sam2.create_job_from_annotations(project_id, file_id, email_id) + + # Verify the request was made with correct URL and payload + mock_client.make_request.assert_called_once() + args, kwargs = mock_client.make_request.call_args + + assert args[0] == "POST" + assert "create_job_from_annotations" in args[1] + assert "client_id=test_client_id" in args[1] + + payload = kwargs["json"] + assert payload["project_id"] == project_id + assert payload["file_id"] == file_id + assert payload["email_id"] == email_id + + assert result["job_ids"] == ["job_123"]