diff --git a/supervision/classification/core.py b/supervision/classification/core.py index c45db8bb..fe5a747f 100644 --- a/supervision/classification/core.py +++ b/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") diff --git a/supervision/dataset/core.py b/supervision/dataset/core.py index 167b6862..76f74006 100644 --- a/supervision/dataset/core.py +++ b/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) diff --git a/test/classification/test_core.py b/test/classification/test_core.py index e742720d..f00c44b0 100644 --- a/test/classification/test_core.py +++ b/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])