add test cases for get_top_k, add docs for Classifications and ClassificationsDataset

docs/classification/core.md

@@ -0,0 +1,3 @@
+## Classifications
+
+:::supervision.classification.core.Classifications

docs/dataset/core.md

@@ -5,4 +5,8 @@
 
 ## DetectionDataset
 
-:::supervision.dataset.core.DetectionDataset
+:::supervision.dataset.core.DetectionDataset
+
+## ClassificationDataset
+
+:::supervision.dataset.core.ClassificationDataset

supervision/__init__.py

@@ -1,7 +1,11 @@
 __version__ = "0.9.0"
 
 from supervision.classification.core import Classifications
-from supervision.dataset.core import BaseDataset, DetectionDataset, ClassificationDataset
+from supervision.dataset.core import (
+    BaseDataset,
+    ClassificationDataset,
+    DetectionDataset,
+)
 from supervision.detection.annotate import BoxAnnotator, MaskAnnotator
 from supervision.detection.core import Detections
 from supervision.detection.line_counter import LineZone, LineZoneAnnotator

supervision/dataset/core.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from pathlib import Path
@@ -8,9 +9,7 @@ from typing import Dict, Iterator, List, Optional, Tuple
 import cv2
 import numpy as np
 
-import os
 from supervision.classification.core import Classifications
-
 from supervision.dataset.formats.pascal_voc import (
     detections_to_pascal_voc,
     load_pascal_voc_annotations,

test/classification/test_core.py

@@ -0,0 +1,39 @@
+from typing import List, TypeVar, Optional, Tuple
+from contextlib import ExitStack as DoesNotRaise
+
+import pytest
+import numpy as np
+
+from supervision.classification.core import Classifications
+
+T = TypeVar("T")
+
+
+@pytest.mark.parametrize(
+    'class_id, confidence, expected_result, exception',
+    [
+        (
+            [0, 1, 2, 3, 4],
+            [0.1, 0.2, 0.9, 0.4, 0.5],
+            (np.array([2, 4, 3, 1, 0]), np.array([0.9, 0.5, 0.4, 0.2, 0.1])),
+            DoesNotRaise()
+        ),  # class_id with 5 numbers and 5 confidences
+        (
+            [0, 1, 2, 3, 4],
+            [0.1, 0.2, 0.3, 0.4],
+            None,
+            pytest.raises(ValueError)
+        ),  # class_id with 5 numbers and 4 confidences
+    ]
+)
+def test_top_k(
+    class_id: List[T],
+    confidence: Optional[List[T]],
+    expected_result: Optional[Tuple[List[T], List[T]]],
+    exception: Exception
+) -> None:
+    with exception:
+        result = Classifications(class_id=np.array(class_id), confidence=np.array(confidence)).get_top_k(len(class_id))
+
+        assert result[0].tolist() == expected_result[0].tolist()
+        assert result[1].tolist() == expected_result[1].tolist()