Repository: roboflow/supervision
File: supervision/dataset/core.py (Python)

Code before:
            )
        if data_yaml_path is not None:
            save_data_yaml(data_yaml_path=data_yaml_path, classes=self.classes)


@dataclass
class ClassificationDataset(BaseDataset):
    """
    Dataclass containing information about a classification dataset.

    Attributes:
        classes (List[str]): List containing dataset class names.
        images (Dict[str, np.ndarray]): Dictionary mapping image name to image.
        annotations (Dict[str, Detections]): Dictionary mapping image name to annotations.
    """

    classes: List[str]
    images: Dict[str, np.ndarray]
    annotations: Dict[str, Classifications]

    def __len__(self) -> int:
        return len(self.images)

    def split(
        self, split_ratio=0.8, random_state=None, shuffle: bool = True
    ) -> Tuple["ClassificationDataset", "ClassificationDataset"]:
        """
        Splits the dataset into two parts (training and testing) using the provided split_ratio.

        Args:
            split_ratio (float, optional): The ratio of the training set to the entire dataset. Default is 0.8.
            random_state (int, optional): The seed for the random number generator. This is used for reproducibility. Default is None.
            shuffle (bool, optional): Whether to shuffle the data before splitting. Default is True.

        Returns:
            Tuple[ClassificationDataset, ClassificationDataset]: A tuple containing the training and testing datasets.

        Example:
            ```python
            >>> import supervision as sv

            >>> cd = sv.ClassificationDataset(...)
            >>> train_cd, test_cd = cd.split(split_ratio=0.7, random_state=42, shuffle=True)
            >>> len(train_cd), len(test_cd)
            (700, 300)
            ```
        """
        image_names = list(self.images.keys())
        train_names, test_names = train_test_split(
            data=image_names,
            train_ratio=split_ratio,

Review comment (suggestion):
Because we use `from __future__ import annotations` at the top of the file, we do not need `"` around `"ClassificationDataset"`. We can just have `Tuple[ClassificationDataset, ClassificationDataset]`.

Code after:
            )
        if data_yaml_path is not None:
            save_data_yaml(data_yaml_path=data_yaml_path, classes=self.classes)


@dataclass
class ClassificationDataset(BaseDataset):
    """
    Dataclass containing information about a classification dataset.

    Attributes:
        classes (List[str]): List containing dataset class names.
        images (Dict[str, np.ndarray]): Dictionary mapping image name to image.
        annotations (Dict[str, Detections]): Dictionary mapping image name to annotations.
    """

    classes: List[str]
    images: Dict[str, np.ndarray]
    annotations: Dict[str, Classifications]

    def __len__(self) -> int:
        return len(self.images)

    def split(
        self, split_ratio=0.8, random_state=None, shuffle: bool = True
    ) -> Tuple[ClassificationDataset, ClassificationDataset]:
        """
        Splits the dataset into two parts (training and testing) using the provided split_ratio.

        Args:
            split_ratio (float, optional): The ratio of the training set to the entire dataset.
            random_state (int, optional): The seed for the random number generator. This is used for reproducibility.
            shuffle (bool, optional): Whether to shuffle the data before splitting.

        Returns:
            Tuple[ClassificationDataset, ClassificationDataset]: A tuple containing the training and testing datasets.

        Example:
            ```python
            >>> import supervision as sv

            >>> cd = sv.ClassificationDataset(...)
            >>> train_cd, test_cd = cd.split(split_ratio=0.7, random_state=42, shuffle=True)
            >>> len(train_cd), len(test_cd)
            (700, 300)
            ```
        """
        image_names = list(self.images.keys())
        train_names, test_names = train_test_split(
            data=image_names,
            train_ratio=split_ratio,