jigsawstack-python/tests/test_object_detection.py at 65902c47ff57bfecbd3f9a505199db674411cfa4 · JigsawStack/jigsawstack-python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
import logging
import os

import pytest
import requests
from dotenv import load_dotenv

import jigsawstack
from jigsawstack.exceptions import JigsawStackError

load_dotenv()


logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


jigsaw = jigsawstack.JigsawStack(api_key=os.getenv("JIGSAWSTACK_API_KEY"))
async_jigsaw = jigsawstack.AsyncJigsawStack(api_key=os.getenv("JIGSAWSTACK_API_KEY"))

IMAGE_URL = (
    "https://rogilvkqloanxtvjfrkm.supabase.co/storage/v1/object/public/demo/Collabo%201080x842.jpg"
)

TEST_CASES = [
    {
        "name": "with_url_only",
        "params": {"url": IMAGE_URL},
        "blob": None,
        "options": None,
    },
    {
        "name": "with_blob_only",
        "params": None,
        "blob": IMAGE_URL,
        "options": None,
    },
    {
        "name": "annotated_image_true",
        "blob": IMAGE_URL,
        "options": {"annotated_image": True},
    },
    {
        "name": "with_annotated_image_false",
        "blob": IMAGE_URL,
        "options": {"annotated_image": False},
    },
    {
        "name": "with_blob_both_features",
        "blob": IMAGE_URL,
        "options": {
            "features": ["object", "gui"],
            "annotated_image": True,
            "return_type": "url",
        },
    },
    {
        "name": "with_blob_gui_features",
        "blob": IMAGE_URL,
        "options": {"features": ["gui"], "annotated_image": False},
    },
    {
        "name": "with_blob_object_detection_features",
        "blob": IMAGE_URL,
        "options": {
            "features": ["object"],
            "annotated_image": True,
            "return_type": "base64",
        },
    },
    {
        "name": "with_prompts",
        "blob": IMAGE_URL,
        "options": {
            "prompts": ["castle", "tree"],
            "annotated_image": True,
        },
    },
    {
        "name": "with_all_options",
        "blob": IMAGE_URL,
        "options": {
            "features": ["object", "gui"],
            "prompts": ["car", "road", "tree"],
            "annotated_image": True,
            "return_type": "base64",
            "return_masks": False,
        },
    },
]


class TestObjectDetectionSync:
    """Test synchronous object detection methods"""

    sync_test_cases = TEST_CASES

    @pytest.mark.parametrize(
        "test_case", sync_test_cases, ids=[tc["name"] for tc in sync_test_cases]
    )
    def test_object_detection(self, test_case):
        """Test synchronous object detection with various inputs"""
        try:
            if test_case.get("blob"):
                # Download blob content
                blob_content = requests.get(test_case["blob"]).content
                result = jigsaw.vision.object_detection(blob_content, test_case.get("options", {}))
            else:
                # Use params directly
                result = jigsaw.vision.object_detection(test_case["params"])

            print(f"Test {test_case['name']}: {result}")
            assert result["success"]
        except JigsawStackError as e:
            pytest.fail(f"Unexpected JigsawStackError in {test_case['name']}: {e}")


class TestObjectDetectionAsync:
    """Test asynchronous object detection methods"""

    async_test_cases = TEST_CASES

    @pytest.mark.parametrize(
        "test_case", async_test_cases, ids=[tc["name"] for tc in async_test_cases]
    )
    @pytest.mark.asyncio
    async def test_object_detection_async(self, test_case):
        """Test asynchronous object detection with various inputs"""
        try:
            if test_case.get("blob"):
                # Download blob content
                blob_content = requests.get(test_case["blob"]).content
                result = await async_jigsaw.vision.object_detection(
                    blob_content, test_case.get("options", {})
                )
            else:
                # Use params directly
                result = await async_jigsaw.vision.object_detection(test_case["params"])

            print(f"Test {test_case['name']}: {result}")
            assert result["success"]
        except JigsawStackError as e:
            pytest.fail(f"Unexpected JigsawStackError in {test_case['name']}: {e}")