dragoneye-python 0.4.1__tar.gz → 1.0.0__tar.gz

{dragoneye_python-0.4.1 → dragoneye_python-1.0.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dragoneye-python
-Version: 0.4.1
+Version: 1.0.0
 License-Expression: MIT
 Requires-Python: >=3.8
 Requires-Dist: requests
@@ -8,3 +8,4 @@ Requires-Dist: pydantic>=2
 Requires-Dist: typing-extensions>=4.0.0
 Requires-Dist: backoff>=2.0.0
 Requires-Dist: aiohttp
+Requires-Dist: polars>=1.0.0

{dragoneye_python-0.4.1 → dragoneye_python-1.0.0}/dragoneye_python.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dragoneye-python
-Version: 0.4.1
+Version: 1.0.0
 License-Expression: MIT
 Requires-Python: >=3.8
 Requires-Dist: requests
@@ -8,3 +8,4 @@ Requires-Dist: pydantic>=2
 Requires-Dist: typing-extensions>=4.0.0
 Requires-Dist: backoff>=2.0.0
 Requires-Dist: aiohttp
+Requires-Dist: polars>=1.0.0

{dragoneye_python-0.4.1 → dragoneye_python-1.0.0}/dragoneye_python.egg-info/SOURCES.txt

@@ -11,6 +11,7 @@ src/dragoneye/classification.py
 src/dragoneye/client.py
 src/dragoneye/constants.py
 src/dragoneye/models.py
+src/dragoneye/parquet_deserializer.py
 src/dragoneye/types/__init__.py
 src/dragoneye/types/common.py
 src/dragoneye/types/exception.py

dragoneye_python-0.4.1/requirements.txt → dragoneye_python-1.0.0/dragoneye_python.egg-info/requires.txt

@@ -2,4 +2,5 @@ requests
 pydantic>=2
 typing-extensions>=4.0.0
 backoff>=2.0.0
-aiohttp
+aiohttp
+polars>=1.0.0

{dragoneye_python-0.4.1 → dragoneye_python-1.0.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "dragoneye-python"
-version = "0.4.1"
+version = "1.0.0"
 requires-python = ">=3.8"
 dynamic = ["dependencies"]
 license = "MIT"

dragoneye_python-0.4.1/dragoneye_python.egg-info/requires.txt → dragoneye_python-1.0.0/requirements.txt

@@ -3,3 +3,4 @@ pydantic>=2
 typing-extensions>=4.0.0
 backoff>=2.0.0
 aiohttp
+polars>=1.0.0

dragoneye_python-1.0.0/src/dragoneye/__init__.py

@@ -0,0 +1,44 @@
+from .classification import (
+    Classification,
+)
+from .client import Dragoneye
+from .models import (
+    ClassificationAttributeOption,
+    ClassificationAttributeResponse,
+    ClassificationCategory,
+    ClassificationCategoryPrediction,
+    ClassificationObjectPrediction,
+    ClassificationPredictImageResponse,
+    ClassificationPredictVideoResponse,
+    ClassificationVideoObjectPrediction,
+)
+from .types.common import NormalizedBbox
+from .types.exception import (
+    PredictionTaskBeginError,
+    PredictionTaskError,
+    PredictionTaskResultsUnavailableError,
+    PredictionTimeoutException,
+    PredictionUploadError,
+)
+from .types.media import Image, Video
+
+__all__ = [
+    "Classification",
+    "ClassificationAttributeOption",
+    "ClassificationAttributeResponse",
+    "ClassificationCategory",
+    "ClassificationCategoryPrediction",
+    "ClassificationObjectPrediction",
+    "ClassificationPredictImageResponse",
+    "ClassificationPredictVideoResponse",
+    "ClassificationVideoObjectPrediction",
+    "Dragoneye",
+    "Image",
+    "NormalizedBbox",
+    "PredictionTaskBeginError",
+    "PredictionTaskError",
+    "PredictionTaskResultsUnavailableError",
+    "PredictionTimeoutException",
+    "PredictionUploadError",
+    "Video",
+]

{dragoneye_python-0.4.1 → dragoneye_python-1.0.0}/src/dragoneye/classification.py

@@ -2,10 +2,12 @@ import asyncio
 import logging
 import time
 from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, overload
+from urllib.parse import urlencode
 
 import aiohttp
 import backoff
 from aiohttp import ClientError
+from multidict import CIMultiDict, CIMultiDictProxy
 from pydantic import BaseModel
 
 from .constants import FAILED_STATUS_PREFIX, PREDICTED_STATUS
@@ -14,6 +16,10 @@ from .models import (
     ClassificationPredictVideoResponse,
     PredictionTaskStatusResponse,
 )
+from .parquet_deserializer import (
+    deserialize_image_predictions,
+    deserialize_video_predictions,
+)
 from .types.common import (
     BASE_API_URL,
     PredictionTaskState,
@@ -129,8 +135,9 @@ class Classification:
         """
         Given a prediction task UUID, return
         """
-        url = f"{BASE_API_URL}/prediction-task/status?predictionTaskUuid={prediction_task_uuid}"
-        headers = {"Authorization": f"Bearer {self._client.api_key}"}
+        query = urlencode({"predictionTaskUuid": prediction_task_uuid})
+        url = f"{BASE_API_URL}/prediction-task/status?{query}"
+        headers = {"X-API-Key": self._client.api_key}
 
         @self._backoff_on_429
         async def _make_request():
@@ -187,31 +194,60 @@ class Classification:
     async def _get_results_unified(
         self, prediction_task_uuid: PredictionTaskUUID, prediction_type: PredictionType
     ) -> Union[ClassificationPredictImageResponse, ClassificationPredictVideoResponse]:
-        url = f"{BASE_API_URL}/prediction-task/results?predictionTaskUuid={prediction_task_uuid}"
-        headers = {"Authorization": f"Bearer {self._client.api_key}"}
+        query = urlencode(
+            {
+                "predictionTaskUuid": prediction_task_uuid,
+                "response_version": "parquet",
+            }
+        )
+        url = f"{BASE_API_URL}/prediction-task/results?{query}"
+        headers = {"X-API-Key": self._client.api_key}
 
         @self._backoff_on_429
-        async def _make_request():
+        async def _make_request() -> tuple[bytes, CIMultiDictProxy[str]]:
             async with aiohttp.ClientSession() as session:
                 async with session.get(url, headers=headers) as resp:
+                    if resp.status == 400:
+                        payload = await resp.json()
+                        raise PredictionTaskResultsUnavailableError(
+                            payload.get("detail", "")
+                        )
                     resp.raise_for_status()
-                    payload = await resp.json()
-            return payload
+                    body_bytes = await resp.read()
+                    response_headers = CIMultiDictProxy(CIMultiDict(resp.headers))
+            return body_bytes, response_headers
 
         try:
-            payload = await _make_request()
+            parquet_bytes, response_headers = await _make_request()
+        except PredictionTaskResultsUnavailableError:
+            raise
         except ClientError as error:
             raise PredictionTaskResultsUnavailableError(
                 f"Error getting prediction task results: {error}"
             )
 
-        # Add the prediction task uuid to the response before returning
-        payload["prediction_task_uuid"] = prediction_task_uuid
+        original_file_name = response_headers.get("X-Original-File-Name")
 
         if prediction_type == "image":
-            return ClassificationPredictImageResponse.model_validate(payload)
+            return ClassificationPredictImageResponse(
+                object_predictions=deserialize_image_predictions(parquet_bytes),
+                prediction_task_uuid=prediction_task_uuid,
+                original_file_name=original_file_name,
+            )
         elif prediction_type == "video":
-            return ClassificationPredictVideoResponse.model_validate(payload)
+            frames_per_second_header = response_headers.get("X-Frames-Per-Second")
+            if frames_per_second_header is None:
+                raise PredictionTaskResultsUnavailableError(
+                    "Missing X-Frames-Per-Second header on video prediction response"
+                )
+            return ClassificationPredictVideoResponse(
+                timestamp_us_to_predictions=deserialize_video_predictions(
+                    parquet_bytes
+                ),
+                frames_per_second=int(frames_per_second_header),
+                prediction_task_uuid=prediction_task_uuid,
+                original_file_name=original_file_name,
+            )
         else:
             raise ValueError(f"Unsupported prediction type: {prediction_type}")
 
@@ -255,13 +291,13 @@ class Classification:
         )
 
         predict_url = f"{BASE_API_URL}/predict"
-        predict_data = {
+        predict_data: dict[str, Any] = {
             "model_name": model_name,
             "prediction_task_uuid": prediction_task_begin_response.prediction_task_uuid,
             **kwargs,
         }
         predict_headers = {
-            "Authorization": f"Bearer {self._client.api_key}",
+            "X-API-Key": self._client.api_key,
         }
 
         @self._backoff_on_429
@@ -354,6 +390,7 @@ class Classification:
 
         form_data = aiohttp.FormData()
         form_data.add_field("mimetype", mime_type)
+        form_data.add_field("response_version", "parquet")
         if file_name is not None:
             form_data.add_field("file_name", file_name)
 
@@ -361,19 +398,24 @@ class Classification:
             form_data.add_field("frames_per_second", str(frames_per_second))
 
         headers = {
-            "Authorization": f"Bearer {self._client.api_key}",
+            "X-API-Key": self._client.api_key,
         }
 
         @self._backoff_on_429
         async def _make_request():
             async with aiohttp.ClientSession() as session:
                 async with session.post(url, data=form_data, headers=headers) as resp:
+                    if resp.status == 400:
+                        error_payload = await resp.json()
+                        raise PredictionTaskBeginError(error_payload.get("detail", ""))
                     resp.raise_for_status()
                     payload = await resp.json()
             return payload
 
         try:
             payload = await _make_request()
+        except PredictionTaskBeginError:
+            raise
         except ClientError as error:
             raise PredictionTaskBeginError("Error beginning prediction task:", error)
 

{dragoneye_python-0.4.1 → dragoneye_python-1.0.0}/src/dragoneye/models.py

@@ -1,4 +1,4 @@
-from typing import Dict, Sequence
+from typing import Dict, List, Optional
 
 from pydantic import BaseModel
 
@@ -7,8 +7,6 @@ from dragoneye.types.common import (
     PredictionTaskState,
     PredictionTaskUUID,
     PredictionType,
-    TaxonID,
-    TaxonPrediction,
 )
 
 
@@ -18,22 +16,38 @@ class PredictionTaskStatusResponse(BaseModel):
     status: PredictionTaskState
 
 
-class ClassificationTraitRootPrediction(BaseModel):
-    id: TaxonID
+class ClassificationAttributeOption(BaseModel):
+    option_id: int
     name: str
-    displayName: str
-    taxons: Sequence[TaxonPrediction]
+    score: float
+
+
+class ClassificationAttributeResponse(BaseModel):
+    attribute_id: int
+    name: str
+    options: List[ClassificationAttributeOption]
+
+
+class ClassificationCategory(BaseModel):
+    id: int
+    name: str
+    score: float
+
+
+class ClassificationCategoryPrediction(BaseModel):
+    category: ClassificationCategory
+    attributes: List[ClassificationAttributeResponse]
 
 
 class ClassificationObjectPrediction(BaseModel):
     normalizedBbox: NormalizedBbox
-    category: TaxonPrediction
-    traits: Sequence[ClassificationTraitRootPrediction]
+    predictions: List[ClassificationCategoryPrediction]
 
 
 class ClassificationPredictImageResponse(BaseModel):
-    predictions: Sequence[ClassificationObjectPrediction]
+    object_predictions: List[ClassificationObjectPrediction]
     prediction_task_uuid: PredictionTaskUUID
+    original_file_name: Optional[str]
 
 
 class ClassificationVideoObjectPrediction(ClassificationObjectPrediction):
@@ -43,7 +57,8 @@ class ClassificationVideoObjectPrediction(ClassificationObjectPrediction):
 
 class ClassificationPredictVideoResponse(BaseModel):
     timestamp_us_to_predictions: Dict[
-        int, Sequence[ClassificationVideoObjectPrediction]
+        int, List[ClassificationVideoObjectPrediction]
     ]
     frames_per_second: int
     prediction_task_uuid: PredictionTaskUUID
+    original_file_name: Optional[str]

dragoneye_python-1.0.0/src/dragoneye/parquet_deserializer.py

@@ -0,0 +1,136 @@
+"""Convert a zstd-compressed parquet prediction blob into typed SDK models.
+
+Matches the server-side schema produced by the dragoneye pipeline:
+
+    image_id:             String
+    normalized_bbox:      Array(Float64, 4)
+    bbox_score:           Float64
+    predictions:          List(Struct{
+                              category_id: Int64,
+                              name:        String,
+                              score:       Float64,
+                              attributes:  List(Struct{
+                                  attribute_id: Int64,
+                                  name:         String,
+                                  options:      List(Struct{
+                                      option_id: Int64,
+                                      name:      String,
+                                      score:     Float64,
+                                  }),
+                              }),
+                          })
+
+For video parquet blobs the extra column ``timestamp_microseconds: Int64`` is
+present.
+"""
+
+import io
+from typing import Any, Dict, List
+
+import polars as pl
+
+from .models import (
+    ClassificationAttributeOption,
+    ClassificationAttributeResponse,
+    ClassificationCategory,
+    ClassificationCategoryPrediction,
+    ClassificationObjectPrediction,
+    ClassificationVideoObjectPrediction,
+)
+from .types.common import NormalizedBbox
+
+
+def _predictions_to_models(
+    raw_predictions: List[Dict[str, Any]],
+) -> List[ClassificationCategoryPrediction]:
+    return [
+        ClassificationCategoryPrediction(
+            category=ClassificationCategory(
+                id=pred["category_id"],
+                name=pred["name"],
+                score=pred["score"],
+            ),
+            attributes=[
+                ClassificationAttributeResponse(
+                    attribute_id=attr["attribute_id"],
+                    name=attr["name"],
+                    options=[
+                        ClassificationAttributeOption(
+                            option_id=opt["option_id"],
+                            name=opt["name"],
+                            score=opt["score"],
+                        )
+                        for opt in attr["options"]
+                    ],
+                )
+                for attr in pred["attributes"]
+            ],
+        )
+        for pred in raw_predictions
+    ]
+
+
+def _read_dataframe(parquet_bytes: bytes, columns: List[str]) -> pl.DataFrame:
+    return pl.read_parquet(io.BytesIO(parquet_bytes), columns=columns)
+
+
+def deserialize_image_predictions(
+    parquet_bytes: bytes,
+) -> List[ClassificationObjectPrediction]:
+    df = _read_dataframe(
+        parquet_bytes, columns=["normalized_bbox", "predictions"]
+    ).filter(pl.col("normalized_bbox").is_not_null())
+    return [
+        ClassificationObjectPrediction(
+            normalizedBbox=NormalizedBbox(tuple(row["normalized_bbox"])),
+            predictions=_predictions_to_models(row["predictions"] or []),
+        )
+        for row in df.select("normalized_bbox", "predictions").iter_rows(named=True)
+    ]
+
+
+def deserialize_video_predictions(
+    parquet_bytes: bytes,
+) -> Dict[int, List[ClassificationVideoObjectPrediction]]:
+    df = _read_dataframe(
+        parquet_bytes,
+        columns=[
+            "image_id",
+            "normalized_bbox",
+            "predictions",
+            "timestamp_microseconds",
+        ],
+    )
+
+    # Frames with detections contribute rows; frames without still appear with a
+    # null bbox so the caller can distinguish "no detections" from "frame not
+    # processed".
+    real_df = df.filter(pl.col("normalized_bbox").is_not_null())
+
+    result: Dict[int, List[ClassificationVideoObjectPrediction]] = {
+        int(ts): [
+            ClassificationVideoObjectPrediction(
+                normalizedBbox=NormalizedBbox(tuple(row["normalized_bbox"])),
+                predictions=_predictions_to_models(row["predictions"] or []),
+                frame_id=row["image_id"],
+                timestamp_microseconds=int(ts),
+            )
+            for row in group.select(
+                "normalized_bbox", "predictions", "image_id"
+            ).iter_rows(named=True)
+        ]
+        for (ts,), group in real_df.group_by(
+            "timestamp_microseconds",
+        )
+    }
+
+    # Include no-detection frames with empty prediction lists.
+    if len(df):
+        for row in df.select("timestamp_microseconds").unique().iter_rows(named=True):
+            ts_raw = row["timestamp_microseconds"]
+            if ts_raw is None:
+                continue
+            ts = int(ts_raw)
+            result.setdefault(ts, [])
+
+    return result

dragoneye_python-1.0.0/src/dragoneye/types/common.py

@@ -0,0 +1,12 @@
+from typing import Literal, NewType, Tuple
+
+PredictionType = Literal["image", "video"]
+PredictionTaskState = NewType("PredictionTaskState", str)
+
+NormalizedBbox = NewType("NormalizedBbox", Tuple[float, float, float, float])
+
+PredictionTaskUUID = NewType("PredictionTaskUUID", str)
+
+TimestampUs = NewType("TimestampUs", int)
+
+BASE_API_URL = "https://api.dragoneye.ai"

{dragoneye_python-0.4.1 → dragoneye_python-1.0.0}/src/dragoneye/types/media.py

@@ -136,10 +136,10 @@ class Media:
 
         # For any readable object with .read()
         if hasattr(src, "read"):
-            pos = _tell_safe(src)
-            data = src.read()
-            _seek_safe(src, pos)
-            return BytesIO(data)
+            pos = _tell_safe(src) # pyright: ignore [reportArgumentType]
+            data = src.read() # pyright: ignore [reportUnknownMemberType, reportAttributeAccessIssue, reportUnknownVariableType]
+            _seek_safe(src, pos) # pyright: ignore [reportArgumentType]
+            return BytesIO(data) # pyright: ignore [reportUnknownArgumentType]
 
         raise TypeError(
             "Invalid media source: expected bytes, BytesIO, or a readable binary stream."

dragoneye_python-0.4.1/src/dragoneye/__init__.py

@@ -1,25 +0,0 @@
-from .classification import (
-    Classification,
-)
-from .client import Dragoneye
-from .models import (
-    ClassificationObjectPrediction,
-    ClassificationPredictImageResponse,
-    ClassificationTraitRootPrediction,
-)
-from .types.common import NormalizedBbox, TaxonID, TaxonPrediction, TaxonType
-from .types.media import Image, Video
-
-__all__ = [
-    "Classification",
-    "ClassificationObjectPrediction",
-    "ClassificationPredictImageResponse",
-    "ClassificationTraitRootPrediction",
-    "Dragoneye",
-    "Image",
-    "Video",
-    "NormalizedBbox",
-    "TaxonID",
-    "TaxonPrediction",
-    "TaxonType",
-]

dragoneye_python-0.4.1/src/dragoneye/types/common.py

@@ -1,31 +0,0 @@
-from enum import Enum
-from typing import Literal, NewType, Optional, Sequence, Tuple
-
-from pydantic import BaseModel
-
-PredictionType = Literal["image", "video"]
-PredictionTaskState = NewType("PredictionTaskState", str)
-
-NormalizedBbox = NewType("NormalizedBbox", Tuple[float, float, float, float])
-
-
-class TaxonType(str, Enum):
-    CATEGORY = ("category",)
-    TRAIT = ("trait",)
-
-
-TaxonID = NewType("TaxonID", int)
-
-PredictionTaskUUID = NewType("PredictionTaskUUID", str)
-
-
-class TaxonPrediction(BaseModel):
-    id: TaxonID
-    type: TaxonType
-    name: str
-    displayName: str
-    score: Optional[float]
-    children: Sequence["TaxonPrediction"]
-
-
-BASE_API_URL = "https://api.dragoneye.ai"