✨ add support for Classification

Author: sebastianMindee

mindee/v2/__init__.py

@@ -1,3 +1,9 @@
+from mindee.v2.product.classification.classification_parameters import (
+    ClassificationParameters,
+)
+from mindee.v2.product.classification.classification_response import (
+    ClassificationResponse,
+)
 from mindee.v2.product.crop.crop_parameters import CropParameters
 from mindee.v2.product.crop.crop_response import CropResponse
 from mindee.v2.product.ocr.ocr_parameters import OCRParameters
@@ -6,6 +12,8 @@
 from mindee.v2.product.split.split_response import SplitResponse
 
 __all__ = [
+    "ClassificationResponse",
+    "ClassificationParameters",
     "CropResponse",
     "CropParameters",
     "OCRResponse",

mindee/v2/product/__init__.py

@@ -1,3 +1,9 @@
+from mindee.v2.product.classification.classification_parameters import (
+    ClassificationParameters,
+)
+from mindee.v2.product.classification.classification_response import (
+    ClassificationResponse,
+)
 from mindee.v2.product.crop.crop_parameters import CropParameters
 from mindee.v2.product.crop.crop_response import CropResponse
 from mindee.v2.product.ocr.ocr_parameters import OCRParameters
@@ -6,6 +12,8 @@
 from mindee.v2.product.split.split_response import SplitResponse
 
 __all__ = [
+    "ClassificationParameters",
+    "ClassificationResponse",
     "CropResponse",
     "CropParameters",
     "OCRResponse",

mindee/v2/product/classification/__init__.py

@@ -0,0 +1,21 @@
+from mindee.v2.product.classification.classification_classifier import (
+    ClassificationClassifier,
+)
+from mindee.v2.product.classification.classification_inference import (
+    ClassificationInference,
+)
+from mindee.v2.product.classification.classification_parameters import (
+    ClassificationParameters,
+)
+from mindee.v2.product.classification.classification_response import (
+    ClassificationResponse,
+)
+from mindee.v2.product.classification.classification_result import ClassificationResult
+
+__all__ = [
+    "ClassificationClassifier",
+    "ClassificationInference",
+    "ClassificationParameters",
+    "ClassificationResponse",
+    "ClassificationResult",
+]

mindee/v2/product/classification/classification_classifier.py

@@ -0,0 +1,14 @@
+from mindee.parsing.common.string_dict import StringDict
+
+
+class ClassificationClassifier:
+    """Document level classification."""
+
+    document_type: str
+    """The document type, as identified on given classification values."""
+
+    def __init__(self, server_response: StringDict):
+        self.document_type = server_response["document_type"]
+
+    def __str__(self) -> str:
+        return f":Document Type: {self.document_type}"

mindee/v2/product/classification/classification_inference.py

@@ -0,0 +1,19 @@
+from mindee.parsing.common.string_dict import StringDict
+from mindee.v2.parsing.inference.base_inference import BaseInference
+from mindee.v2.product.classification.classification_result import ClassificationResult
+
+
+class ClassificationInference(BaseInference):
+    """The inference result for a classification utility request."""
+
+    result: ClassificationResult
+    """Result of a classification inference."""
+    _slug: str = "classification"
+    """Slug of the endpoint."""
+
+    def __init__(self, raw_response: StringDict) -> None:
+        super().__init__(raw_response)
+        self.result = ClassificationResult(raw_response["result"])
+
+    def __str__(self) -> str:
+        return f"Inference\n#########\n{self.model}\n{self.file}\n{self.result}\n"

mindee/v2/product/classification/classification_parameters.py

@@ -0,0 +1,9 @@
+from mindee.input.base_parameters import BaseParameters
+
+
+class ClassificationParameters(BaseParameters):
+    """
+    Parameters accepted by the classification utility v2 endpoint.
+    """
+
+    _slug: str = "utilities/classification"

mindee/v2/product/classification/classification_response.py

@@ -0,0 +1,19 @@
+from mindee.parsing.common.string_dict import StringDict
+from mindee.v2.parsing.inference import BaseResponse
+from mindee.v2.product.classification.classification_inference import (
+    ClassificationInference,
+)
+
+
+class ClassificationResponse(BaseResponse):
+    """Represent a classification inference response from Mindee V2 API."""
+
+    inference: ClassificationInference
+    """Inference object for classification inference."""
+
+    _slug: str = "utilities/classification"
+    """Slug of the inference."""
+
+    def __init__(self, raw_response: StringDict) -> None:
+        super().__init__(raw_response)
+        self.inference = ClassificationInference(raw_response["inference"])

mindee/v2/product/classification/classification_result.py

@@ -0,0 +1,16 @@
+from mindee.parsing.common.string_dict import StringDict
+from mindee.v2.product.classification.classification_classifier import (
+    ClassificationClassifier,
+)
+
+
+class ClassificationResult:
+    """Classification result info."""
+
+    classification: ClassificationClassifier
+
+    def __init__(self, raw_response: StringDict) -> None:
+        self.classification = ClassificationClassifier(raw_response["classification"])
+
+    def __str__(self) -> str:
+        return f"Classification\n======{self.classification}"

mindee/v2/product/split/split_range.py

@@ -7,9 +7,12 @@ class SplitRange:
     """Split inference result."""
 
     page_range: List[int]
-    """Page range of the split inference."""
+    """
+    0-based page indexes, where the first integer indicates the start page and the
+    second integer indicates the end page.
+    """
     document_type: str
-    """Document type of the split inference."""
+    """The document type, as identified on given classification values."""
 
     def __init__(self, server_response: StringDict):
         self.page_range = server_response["page_range"]