generated from roboflow/template-python
-
Notifications
You must be signed in to change notification settings - Fork 1.4k
/
utils.py
130 lines (102 loc) Β· 3.8 KB
/
utils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
import copy
import os
import random
from pathlib import Path
from typing import Dict, List, Optional, Tuple, TypeVar
import cv2
import numpy as np
from supervision.detection.core import Detections
from supervision.detection.utils import (
approximate_polygon,
filter_polygons_by_area,
mask_to_polygons,
)
T = TypeVar("T")
def approximate_mask_with_polygons(
mask: np.ndarray,
min_image_area_percentage: float = 0.0,
max_image_area_percentage: float = 1.0,
approximation_percentage: float = 0.75,
) -> List[np.ndarray]:
height, width = mask.shape
image_area = height * width
minimum_detection_area = min_image_area_percentage * image_area
maximum_detection_area = max_image_area_percentage * image_area
polygons = mask_to_polygons(mask=mask)
if len(polygons) == 1:
polygons = filter_polygons_by_area(
polygons=polygons, min_area=None, max_area=maximum_detection_area
)
else:
polygons = filter_polygons_by_area(
polygons=polygons,
min_area=minimum_detection_area,
max_area=maximum_detection_area,
)
return [
approximate_polygon(polygon=polygon, percentage=approximation_percentage)
for polygon in polygons
]
def merge_class_lists(class_lists: List[List[str]]) -> List[str]:
unique_classes = set()
for class_list in class_lists:
for class_name in class_list:
unique_classes.add(class_name.lower())
return sorted(list(unique_classes))
def build_class_index_mapping(
source_classes: List[str], target_classes: List[str]
) -> Dict[int, int]:
index_mapping = {}
for i, class_name in enumerate(source_classes):
if class_name not in target_classes:
raise ValueError(
f"Class {class_name} not found in target classes. "
f"source_classes must be a subset of target_classes."
)
corresponding_index = target_classes.index(class_name)
index_mapping[i] = corresponding_index
return index_mapping
def map_detections_class_id(
source_to_target_mapping: Dict[int, int], detections: Detections
) -> Detections:
if detections.class_id is None:
raise ValueError("Detections must have class_id attribute.")
if set(np.unique(detections.class_id)) - set(source_to_target_mapping.keys()):
raise ValueError(
"Detections class_id must be a subset of source_to_target_mapping keys."
)
detections_copy = copy.deepcopy(detections)
if len(detections) > 0:
detections_copy.class_id = np.vectorize(source_to_target_mapping.get)(
detections_copy.class_id
)
return detections_copy
def save_dataset_images(
images_directory_path: str, images: Dict[str, np.ndarray]
) -> None:
Path(images_directory_path).mkdir(parents=True, exist_ok=True)
for image_name, image in images.items():
target_image_path = os.path.join(images_directory_path, image_name)
cv2.imwrite(target_image_path, image)
def train_test_split(
data: List[T],
train_ratio: float = 0.8,
random_state: Optional[int] = None,
shuffle: bool = True,
) -> Tuple[List[T], List[T]]:
"""
Splits the data into two parts using the provided train_ratio.
Args:
data (List[T]): The data to split.
train_ratio (float): The ratio of the training set to the entire dataset.
random_state (Optional[int]): The seed for the random number generator.
shuffle (bool): Whether to shuffle the data before splitting.
Returns:
Tuple[List[T], List[T]]: The split data.
"""
if random_state is not None:
random.seed(random_state)
if shuffle:
random.shuffle(data)
split_index = int(len(data) * train_ratio)
return data[:split_index], data[split_index:]