Skip to content

Commit 6b0be9c

Browse files
SkalskiPSkalskiP
authored
Merge pull request roboflow#296 from roboflow/fix/yoloexplorer
Image Path as keys for DetectionDataset
2 parents 6f03b7f + cf2c8cb commit 6b0be9c

5 files changed

Lines changed: 44 additions & 39 deletions

File tree

‎supervision/dataset/core.py‎

Lines changed: 19 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -181,21 +181,21 @@ def as_pascal_voc(
181181
in the range [0, 1). Argument is used only for segmentation datasets.
182182
"""
183183
if images_directory_path:
184-
images_path = Path(images_directory_path)
185-
images_path.mkdir(parents=True, exist_ok=True)
186-
184+
save_dataset_images(
185+
images_directory_path=images_directory_path, images=self.images
186+
)
187187
if annotations_directory_path:
188-
annotations_path = Path(annotations_directory_path)
189-
annotations_path.mkdir(parents=True, exist_ok=True)
188+
Path(annotations_directory_path).mkdir(parents=True, exist_ok=True)
190189

191-
for image_name, image in self.images.items():
192-
detections = self.annotations[image_name]
193-
194-
if images_directory_path:
195-
cv2.imwrite(str(images_path / image_name), image)
190+
for image_path, image in self.images.items():
191+
detections = self.annotations[image_path]
196192

197193
if annotations_directory_path:
198-
annotation_name = Path(image_name).stem
194+
annotation_name = Path(image_path).stem
195+
annotations_path = os.path.join(
196+
annotations_directory_path, f"{annotation_name}.xml"
197+
)
198+
image_name = Path(image_path).name
199199
pascal_voc_xml = detections_to_pascal_voc(
200200
detections=detections,
201201
classes=self.classes,
@@ -206,7 +206,7 @@ def as_pascal_voc(
206206
approximation_percentage=approximation_percentage,
207207
)
208208

209-
with open(annotations_path / f"{annotation_name}.xml", "w") as f:
209+
with open(annotations_path, "w") as f:
210210
f.write(pascal_voc_xml)
211211

212212
@classmethod
@@ -615,9 +615,10 @@ def as_folder_structure(self, root_directory_path: str) -> None:
615615
for class_name in self.classes:
616616
os.makedirs(os.path.join(root_directory_path, class_name), exist_ok=True)
617617

618-
for image_name in self.images:
619-
classification = self.annotations[image_name]
620-
image = self.images[image_name]
618+
for image_path in self.images:
619+
classification = self.annotations[image_path]
620+
image = self.images[image_path]
621+
image_name = Path(image_path).name
621622
class_id = (
622623
classification.class_id[0]
623624
if classification.confidence is None
@@ -666,9 +667,9 @@ def from_folder_structure(cls, root_directory_path: str) -> ClassificationDatase
666667
class_id = classes.index(class_name)
667668

668669
for image in os.listdir(os.path.join(root_directory_path, class_name)):
669-
image_dir = os.path.join(root_directory_path, class_name, image)
670-
images[image] = cv2.imread(image_dir)
671-
annotations[image] = Classifications(
670+
image_path = str(os.path.join(root_directory_path, class_name, image))
671+
images[image_path] = cv2.imread(image_path)
672+
annotations[image_path] = Classifications(
672673
class_id=np.array([class_id]),
673674
)
674675

‎supervision/dataset/formats/coco.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -159,7 +159,7 @@ def load_coco_annotations(
159159
image_annotations = coco_annotations_groups.get(coco_image["id"], [])
160160
image_path = os.path.join(images_directory_path, image_name)
161161

162-
image = cv2.imread(str(image_path))
162+
image = cv2.imread(image_path)
163163
annotation = coco_annotations_to_detections(
164164
image_annotations=image_annotations,
165165
resolution_wh=(image_width, image_height),
@@ -170,8 +170,8 @@ def load_coco_annotations(
170170
detections=annotation,
171171
)
172172

173-
images[image_name] = image
174-
annotations[image_name] = annotation
173+
images[image_path] = image
174+
annotations[image_path] = annotation
175175

176176
return classes, images, annotations
177177

@@ -200,9 +200,9 @@ def save_coco_annotations(
200200
coco_categories = classes_to_coco_categories(classes=classes)
201201

202202
image_id, annotation_id = 1, 1
203-
for image_name, image in images.items():
203+
for image_path, image in images.items():
204204
image_height, image_width, _ = image.shape
205-
205+
image_name = f"{Path(image_path).stem}{Path(image_path).suffix}"
206206
coco_image = {
207207
"id": image_id,
208208
"license": 1,
@@ -213,7 +213,7 @@ def save_coco_annotations(
213213
}
214214

215215
coco_images.append(coco_image)
216-
detections = annotations[image_name]
216+
detections = annotations[image_path]
217217

218218
coco_annotation, annotation_id = detections_to_coco_annotations(
219219
detections=detections,

‎supervision/dataset/formats/pascal_voc.py‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -168,12 +168,13 @@ def load_pascal_voc_annotations(
168168

169169
for image_path in image_paths:
170170
image_name = Path(image_path).stem
171-
image = cv2.imread(str(image_path))
171+
image_path = str(image_path)
172+
image = cv2.imread(image_path)
172173

173174
annotation_path = os.path.join(annotations_directory_path, f"{image_name}.xml")
174175
if not os.path.exists(annotation_path):
175-
images[image_path.name] = image
176-
annotations[image_path.name] = Detections.empty()
176+
images[image_path] = image
177+
annotations[image_path] = Detections.empty()
177178
continue
178179

179180
tree = parse(annotation_path)
@@ -184,8 +185,8 @@ def load_pascal_voc_annotations(
184185
root, classes, resolution_wh, force_masks
185186
)
186187

187-
images[image_path.name] = image
188-
annotations[image_path.name] = annotation
188+
images[image_path] = image
189+
annotations[image_path] = annotation
189190

190191
return classes, images, annotations
191192

‎supervision/dataset/formats/yolo.py‎

Lines changed: 11 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -139,13 +139,14 @@ def load_yolo_annotations(
139139
annotations = {}
140140

141141
for image_path in image_paths:
142-
image_name = Path(image_path).stem
143-
image = cv2.imread(str(image_path))
142+
image_stem = Path(image_path).stem
143+
image_path = str(image_path)
144+
image = cv2.imread(image_path)
144145

145-
annotation_path = os.path.join(annotations_directory_path, f"{image_name}.txt")
146+
annotation_path = os.path.join(annotations_directory_path, f"{image_stem}.txt")
146147
if not os.path.exists(annotation_path):
147-
images[image_path.name] = image
148-
annotations[image_path.name] = Detections.empty()
148+
images[image_path] = image
149+
annotations[image_path] = Detections.empty()
149150
continue
150151

151152
lines = read_txt_file(str(annotation_path))
@@ -158,8 +159,8 @@ def load_yolo_annotations(
158159
lines=lines, resolution_wh=resolution_wh, with_masks=with_masks
159160
)
160161

161-
images[image_path.name] = image
162-
annotations[image_path.name] = annotation
162+
images[image_path] = image
163+
annotations[image_path] = annotation
163164
return classes, images, annotations
164165

165166

@@ -227,8 +228,9 @@ def save_yolo_annotations(
227228
approximation_percentage: float = 0.75,
228229
) -> None:
229230
Path(annotations_directory_path).mkdir(parents=True, exist_ok=True)
230-
for image_name, image in images.items():
231-
detections = annotations[image_name]
231+
for image_path, image in images.items():
232+
detections = annotations[image_path]
233+
image_name = Path(image_path).name
232234
yolo_annotations_name = _image_name_to_annotation_name(image_name=image_name)
233235
yolo_annotations_path = os.path.join(
234236
annotations_directory_path, yolo_annotations_name

‎supervision/dataset/utils.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,8 @@ def save_dataset_images(
9797
) -> None:
9898
Path(images_directory_path).mkdir(parents=True, exist_ok=True)
9999

100-
for image_name, image in images.items():
100+
for image_path, image in images.items():
101+
image_name = Path(image_path).name
101102
target_image_path = os.path.join(images_directory_path, image_name)
102103
cv2.imwrite(target_image_path, image)
103104

0 commit comments

Comments
 (0)