respond to feedback

supervision/classification/core.py

@@ -57,6 +57,7 @@ class Classifications:
             >>> cd.annotations["image.png"].get_top_k(1)
 
             (array([1]), array([0.9]))
+            ```
         """
         if self.confidence is None:
             raise ValueError("confidence is None")

supervision/dataset/core.py

@@ -354,7 +354,7 @@ class ClassificationDataset(BaseDataset):
 
     def split(
         self, split_ratio=0.8, random_state=None, shuffle: bool = True
-    ) -> Tuple["ClassificationDataset", "ClassificationDataset"]:
+    ) -> Tuple[ClassificationDataset, ClassificationDataset]:
         """
         Splits the dataset into two parts (training and testing) using the provided split_ratio.
 
@@ -396,12 +396,12 @@ class ClassificationDataset(BaseDataset):
         )
         return train_dataset, test_dataset
 
-    def as_multiclass_folder_structure(self, output_directory_path: str) -> None:
+    def as_multiclass_folder_structure(self, root_directory_path: str) -> None:
         """
         Saves the dataset as a multi-class folder structure.
 
         Args:
-            output_directory_path (str): The path to the directory where the dataset will be saved.
+            root_directory_path (str): The path to the directory where the dataset will be saved.
 
         Example:
             ```python
@@ -411,26 +411,27 @@ class ClassificationDataset(BaseDataset):
 
             >>> cd.as_multiclass_folder_structure(
             ...     root_directory_path="./images/valid/",
-            ...     output_directory_path="./out",
+            ...     root_directory_path="./out",
             ... )
+            ```
         """
-        os.makedirs(output_directory_path, exist_ok=True)
+        os.makedirs(root_directory_path, exist_ok=True)
 
         for class_name in self.classes:
-            os.makedirs(os.path.join(output_directory_path, class_name), exist_ok=True)
+            os.makedirs(os.path.join(root_directory_path, class_name), exist_ok=True)
 
             for image in self.annotations:
                 for class_id in self.annotations[image].class_id:
                     if class_id == self.classes.index(class_name):
                         cv2.imwrite(
-                            os.path.join(output_directory_path, class_name, image),
+                            os.pafth.join(root_directory_path, class_name, image),
                             self.images[image],
                         )
 
     @classmethod
     def from_multiclass_folder_structure(
         cls, root_directory_path: str
-    ) -> "ClassificationDataset":
+    ) -> ClassificationDataset:
         """
         Load data from a multiclass folder structure into a ClassificationDataset.
 
@@ -447,6 +448,7 @@ class ClassificationDataset(BaseDataset):
             >>> cd = sv.ClassificationDataset.from_multiclass_folder_structure(
             ...     root_directory_path="./dataset",
             ... )
+            ```
         """
         classes = os.listdir(root_directory_path)
 

test/classification/test_core.py

@@ -1,4 +1,4 @@
-from typing import List, TypeVar, Optional, Tuple
+from typing import List, Optional, Tuple
 from contextlib import ExitStack as DoesNotRaise
 
 import pytest
@@ -6,34 +6,58 @@ import numpy as np
 
 from supervision.classification.core import Classifications
 
-T = TypeVar("T")
-
 
 @pytest.mark.parametrize(
-    'class_id, confidence, expected_result, exception',
+    "class_id, confidence, k, expected_result, exception",
     [
         (
-            [0, 1, 2, 3, 4],
-            [0.1, 0.2, 0.9, 0.4, 0.5],
+            np.array([0, 1, 2, 3, 4]),
+            np.array([0.1, 0.2, 0.9, 0.4, 0.5]),
+            5,
             (np.array([2, 4, 3, 1, 0]), np.array([0.9, 0.5, 0.4, 0.2, 0.1])),
-            DoesNotRaise()
+            DoesNotRaise(),
         ),  # class_id with 5 numbers and 5 confidences
+        (
+            np.array([5, 1, 2, 3, 4]),
+            np.array([0.1, 0.2, 0.9, 0.4, 0.5]),
+            1,
+            (np.array([2]), np.array([0.9])),
+            DoesNotRaise(),
+        ),  # class_id with 5 numbers and 5 confidences, retrieve where k = 1
+        (
+            np.array([4, 1, 2, 3, 6, 5]),
+            np.array([0.8, 0.2, 0.9, 0.4, 0.5, 0.1]),
+            2,
+            (np.array([2, 4]), np.array([0.9, 0.8])),
+            DoesNotRaise(),
+        ),  # class_id with 5 numbers and 5 confidences, retrieve where k = 3
+        (
+            np.array([0, 1, 2, 3, 4]),
+            np.array([]),
+            5,
+            None,
+            pytest.raises(ValueError),
+        ),  # class_id with 5 numbers and 0 confidences
         (
             [0, 1, 2, 3, 4],
             [0.1, 0.2, 0.3, 0.4],
+            5,
             None,
-            pytest.raises(ValueError)
+            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
+    class_id: np.ndarray,
+    confidence: Optional[np.ndarray],
+    k: int,
+    expected_result: Optional[Tuple[np.ndarray, np.ndarray]],
+    exception: Exception,
 ) -> None:
     with exception:
-        result = Classifications(class_id=np.array(class_id), confidence=np.array(confidence)).get_top_k(len(class_id))
+        result = Classifications(
+            class_id=np.array(class_id), confidence=np.array(confidence)
+        ).get_top_k(k)
 
-        assert result[0].tolist() == expected_result[0].tolist()
-        assert result[1].tolist() == expected_result[1].tolist()
+        assert np.array_equal(result[0], expected_result[0])
+        assert np.array_equal(result[1], expected_result[1])