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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dragoneye-python
-Version: 0.5.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.5.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.5.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.5.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.5.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.5.1 → dragoneye_python-1.0.0}/pyproject.toml

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

dragoneye_python-0.5.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-0.5.1 → dragoneye_python-1.0.0}/src/dragoneye/__init__.py

@@ -13,6 +13,13 @@ from .models import (
     ClassificationVideoObjectPrediction,
 )
 from .types.common import NormalizedBbox
+from .types.exception import (
+    PredictionTaskBeginError,
+    PredictionTaskError,
+    PredictionTaskResultsUnavailableError,
+    PredictionTimeoutException,
+    PredictionUploadError,
+)
 from .types.media import Image, Video
 
 __all__ = [
@@ -28,5 +35,10 @@ __all__ = [
     "Dragoneye",
     "Image",
     "NormalizedBbox",
+    "PredictionTaskBeginError",
+    "PredictionTaskError",
+    "PredictionTaskResultsUnavailableError",
+    "PredictionTimeoutException",
+    "PredictionUploadError",
     "Video",
 ]

{dragoneye_python-0.5.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}")
 
@@ -261,7 +297,7 @@ class Classification:
             **kwargs,
         }
         predict_headers = {
-            "Authorization": f"Bearer {self._client.api_key}",
+            "X-API-Key": self._client.api_key,
         }
 
         @self._backoff_on_429
@@ -354,7 +390,7 @@ class Classification:
 
         form_data = aiohttp.FormData()
         form_data.add_field("mimetype", mime_type)
-        form_data.add_field("response_version", "object")
+        form_data.add_field("response_version", "parquet")
         if file_name is not None:
             form_data.add_field("file_name", file_name)
 
@@ -362,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-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