okooooooo/supervision
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor code, respond to feedback

Browse Source
This commit is contained in:
James Gallagher committed 2023-06-13 12:48:07 +01:00
1 parent eb1ae1b790
commit b5dedf09cd
2 files changed
+153 -160
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

supervision/classification/core.py
+8 -160
View File
@@ -22,9 +22,10 @@ def _validate_confidence(confidence: Any, n: int) -> None:
"""
Ensure that confidence is a 1d np.ndarray with (n, ) shape.
"""
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")
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
@@ -40,8 +41,7 @@ class Classifications:
_validate_class_ids(self.class_id, n)
if self.confidence is not None:
_validate_confidence(self.confidence, n)
_validate_confidence(self.confidence, n)
def get_top_k(self, k: int) -> Tuple[np.ndarray, np.ndarray]:
"""
@@ -66,161 +66,9 @@ class Classifications:
if self.confidence is None:
raise ValueError("confidence is None")
confidence = self.confidence.copy()
class_ids = self.class_id.copy()
order = np.argsort(confidence.copy())[::-1]
order = np.argsort(self.confidence.copy())[::-1]
top_k_order = order[:k]
top_k_class_id = class_ids[top_k_order]
top_k_confidence = confidence[top_k_order]
top_k_class_id = self.class_ids[top_k_order]
top_k_confidence = self.confidence[top_k_order]
return top_k_class_id, top_k_confidence
@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,
random_state=random_state,
shuffle=shuffle,
)
train_dataset = ClassificationDataset(
classes=self.classes,
images={name: self.images[name] for name in train_names},
annotations={name: self.annotations[name] for name in train_names},
)
test_dataset = ClassificationDataset(
classes=self.classes,
images={name: self.images[name] for name in test_names},
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:
root_directory_path (str): The path to the dataset directory.
Returns:
ClassificationDataset: The dataset.
Example:
```python
>>> import supervision as sv
>>> cd = sv.ClassificationDataset.from_multiclass_folder_structure(
... root_directory_path="./dataset",
... )
"""
classes = os.listdir(root_directory_path)
classes = list(set(classes))
classes.sort()
images = {}
annotations = {}
for class_name in classes:
if not os.path.exists(os.path.join(root_directory_path, class_name)):
continue
for image in os.listdir(os.path.join(root_directory_path, class_name)):
images[image] = cv2.imread(
os.path.join(root_directory_path, class_name, image)
)
annotations[image] = Classifications(
class_id=np.array([classes.index(class_name)]),
)
return cls(
classes=classes,
images=images,
annotations=annotations,
)
supervision/dataset/core.py
+145
View File
@@ -8,6 +8,9 @@ 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,
@@ -330,3 +333,145 @@ class DetectionDataset(BaseDataset):
)
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,
random_state=random_state,
shuffle=shuffle,
)
train_dataset = ClassificationDataset(
classes=self.classes,
images={name: self.images[name] for name in train_names},
annotations={name: self.annotations[name] for name in train_names},
)
test_dataset = ClassificationDataset(
classes=self.classes,
images={name: self.images[name] for name in test_names},
annotations={name: self.annotations[name] for name in test_names},
)
return train_dataset, test_dataset
def as_multiclass_folder_structure(
self, output_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.
Example:
```python
>>> import supervision as sv
>>> cd = sv.ClassificationDataset(...)
>>> cd.as_multiclass_folder_structure(
... root_directory_path="./images/valid/",
... output_directory_path="./out",
... )
"""
os.makedirs(output_directory_path, exist_ok=True)
for class_name in self.classes:
os.makedirs(os.path.join(output_directory_path, class_name), exist_ok=True)
for image in self.annotations:
# sort class_ids
if self.annotations[image].class_id[0] == self.classes.index(
class_name
):
cv2.imwrite(
os.path.join(output_directory_path, class_name, image),
self.images[image],
)
@classmethod
def from_multiclass_folder_structure(
cls, root_directory_path: str
) -> "ClassificationDataset":
"""
Load data from a multiclass folder structure into a ClassificationDataset.
Args:
root_directory_path (str): The path to the dataset directory.
Returns:
ClassificationDataset: The dataset.
Example:
```python
>>> import supervision as sv
>>> cd = sv.ClassificationDataset.from_multiclass_folder_structure(
... root_directory_path="./dataset",
... )
"""
classes = os.listdir(root_directory_path)
classes = sorted(set(classes))
images = {}
annotations = {}
for class_name in classes:
class_id = classes.index(class_name)
for image in os.listdir(os.path.join(root_directory_path, class_name)):
image_dir = os.path.join(root_directory_path, class_name, image)
images[image] = cv2.imread(image_dir)
annotations[image] = Classifications(
class_id=np.array([class_id]),
)
return cls(
classes=classes,
images=images,
annotations=annotations,
)