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

Code before:
            annotations={name: self.annotations[name] for name in test_names},
        )
        return train_dataset, test_dataset

    def as_multiclass_folder_structure(
        self, root_directory_path: str, output_directory_path: str
    ) -> None:
        """
        Saves the dataset as a multi-class folder structure.

        Args:
            root_directory_path (str): The path to the root directory of the images with which you are working.
            output_directory_path (str): The path to the directory where the dataset will be saved.

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

            >>> cd = sv.ClassificationDataset(...)

            >>> cd.as_multiclass_folder_structure(
            ...     root_directory_path="./images/valid/",
            ...     output_directory_path="./out",
            ... )
        """
        if not os.path.exists(output_directory_path):
            os.makedirs(output_directory_path)

        for class_name in self.classes:
            if not os.path.exists(os.path.join(output_directory_path, class_name)):
                os.makedirs(os.path.join(output_directory_path, class_name))

            for image in self.annotations:
                if self.annotations[image].class_id[0] == self.classes.index(
                    class_name
                ):
                    full_dir = os.path.join(root_directory_path, class_name, image)

                    cv2.imwrite(
                        os.path.join(output_directory_path, class_name, image),
                        cv2.imread(full_dir),
                    )

    @classmethod
    def from_multiclass_folder_structure(
        cls, root_directory_path: str
    ) -> "ClassificationDataset":
        """
        Load data from a multiclass folder structure into a ClassificationDataset.

        Args:

Review comment (suggestion):
You don't need to check if the directory exists before creating a directory. You can do: `os.makedirs(path, exist_ok=True)`. It will create dir if it doesn't exist and do nothing if it exists. 

Code after:
from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Optional, Tuple

import numpy as np


def _validate_class_ids(class_id: Any, n: int) -> None:
    """
    Ensure that class_id is a 1d np.ndarray with (n, ) shape.
    """
    is_valid = isinstance(class_id, np.ndarray) and class_id.shape == (n,)
    if not is_valid:
        raise ValueError("class_id must be 1d np.ndarray with (n, ) shape")


def _validate_confidence(confidence: Any, n: int) -> None:
    """
    Ensure that confidence is a 1d np.ndarray with (n, ) shape.
    """
    if confidence is not None:
        is_valid = isinstance(confidence, np.ndarray) and confidence.shape == (n,)
        if not is_valid:
            raise ValueError("confidence must be 1d np.ndarray with (n, ) shape")


@dataclass
class Classifications:
    class_id: np.ndarray
    confidence: Optional[np.ndarray] = None

    def __post_init__(self) -> None:
        """
        Validate the classification inputs.
        """
        n = len(self.class_id)

        _validate_class_ids(self.class_id, n)
        _validate_confidence(self.confidence, n)

    @classmethod
    def from_yolov8(cls, yolov8_results) -> Classifications:
        """
        Creates a Classifications instance from a [YOLOv8](https://github.com/ultralytics/ultralytics) inference result.

        Args:
            yolov8_results (ultralytics.yolo.engine.results.Results): The output Results instance from YOLOv8

        Returns:
            Detections: A new Classifications object.

        Example:
            ```python
            >>> import cv2
            >>> from ultralytics import YOLO
            >>> import supervision as sv

            >>> image = cv2.imread(SOURCE_IMAGE_PATH)
            >>> model = YOLO('yolov8s-cls.pt')
            >>> result = model(image)[0]
            >>> classifications = sv.Classifications.from_yolov8(result)
            ```
        """
        confidence = yolov8_results.probs.data.cpu().numpy()
        return cls(class_id=np.arange(confidence.shape[0]), confidence=confidence)

    def get_top_k(self, k: int) -> Tuple[np.ndarray, np.ndarray]:
        """
        Retrieve the top k class IDs and confidences, ordered in descending order by confidence.

        Args:
            k (int): The number of top class IDs and confidences to retrieve.

        Returns:
            Tuple[np.ndarray, np.ndarray]: A tuple containing the top k class IDs and confidences.

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

            >>> classifications = sv.Classifications(...)

            >>> classifications.get_top_k(1)

            (array([1]), array([0.9]))
            ```
        """
        if self.confidence is None:
            raise ValueError("top_k could not be calculated, confidence is None")

        order = np.argsort(self.confidence)[::-1]
        top_k_order = order[:k]
        top_k_class_id = self.class_id[top_k_order]
        top_k_confidence = self.confidence[top_k_order]

        return top_k_class_id, top_k_confidence