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

small simplifications

This commit is contained in:
SkalskiP 2023-08-22 16:34:54 +02:00
parent 18f7cad081
commit f8baad3065
3 changed files with 20 additions and 24 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

29
supervision/dataset/core.py
View File

@ -182,19 +182,19 @@ class DetectionDataset(BaseDataset):
"""
if images_directory_path:
save_dataset_images(
images_directory_path=images_directory_path, images=self.images
images_directory_path=images_directory_path,
images=self.images
)
if annotations_directory_path:
annotations_path = Path(annotations_directory_path)
annotations_path.mkdir(parents=True, exist_ok=True)
Path(annotations_directory_path).mkdir(parents=True, exist_ok=True)
for image_path, image in self.images.items():
detections = self.annotations[image_path]
if annotations_directory_path:
annotation_name = Path(image_path).stem
image_name = f"{Path(image_path).stem}{Path(image_path).suffix}"
annotations_path = os.path.join(annotations_directory_path, f"{annotation_name}.xml")
image_name = Path(image_path).name
pascal_voc_xml = detections_to_pascal_voc(
detections=detections,
classes=self.classes,
@ -205,7 +205,7 @@ class DetectionDataset(BaseDataset):
approximation_percentage=approximation_percentage,
)
with open(annotations_path / f"{annotation_name}.xml", "w") as f:
with open(annotations_path, "w") as f:
f.write(pascal_voc_xml)
@classmethod
@ -356,7 +356,8 @@ class DetectionDataset(BaseDataset):
"""
if images_directory_path is not None:
save_dataset_images(
images_directory_path=images_directory_path, images=self.images
images_directory_path=images_directory_path,
images=self.images
)
if annotations_directory_path is not None:
save_yolo_annotations(
@ -453,7 +454,8 @@ class DetectionDataset(BaseDataset):
"""
if images_directory_path is not None:
save_dataset_images(
images_directory_path=images_directory_path, images=self.images
images_directory_path=images_directory_path,
images=self.images
)
if annotations_path is not None:
save_coco_annotations(
@ -614,11 +616,10 @@ class ClassificationDataset(BaseDataset):
for class_name in self.classes:
os.makedirs(os.path.join(root_directory_path, class_name), exist_ok=True)
for image_name in self.images:
classification = self.annotations[image_name]
image = self.images[image_name]
image_name = str(Path(image_name).stem)
image_ext = str(Path(image_name).suffix)
for image_path in self.images:
classification = self.annotations[image_path]
image = self.images[image_path]
image_name = Path(image_path).name
class_id = (
classification.class_id[0]
if classification.confidence is None
@ -626,7 +627,7 @@ class ClassificationDataset(BaseDataset):
)
class_name = self.classes[class_id]
image_path = os.path.join(
root_directory_path, class_name, f"{image_name}{image_ext}"
root_directory_path, class_name, image_name
)
cv2.imwrite(image_path, image)

8
supervision/dataset/formats/yolo.py
View File

@ -139,11 +139,11 @@ def load_yolo_annotations(
annotations = {}
for image_path in image_paths:
image_name = Path(image_path).stem
image_stem = Path(image_path).stem
image_path = str(image_path)
image = cv2.imread(image_path)
annotation_path = os.path.join(annotations_directory_path, f"{image_name}.txt")
annotation_path = os.path.join(annotations_directory_path, f"{image_stem}.txt")
if not os.path.exists(annotation_path):
images[image_path] = image
annotations[image_path] = Detections.empty()
@ -230,10 +230,8 @@ def save_yolo_annotations(
Path(annotations_directory_path).mkdir(parents=True, exist_ok=True)
for image_path, image in images.items():
detections = annotations[image_path]
image_name = f"{Path(image_path).stem}{Path(image_path).suffix}"
image_name = Path(image_path).name
yolo_annotations_name = _image_name_to_annotation_name(image_name=image_name)
print(yolo_annotations_name)
yolo_annotations_path = os.path.join(
annotations_directory_path, yolo_annotations_name
)

7
supervision/dataset/utils.py
View File

@ -98,11 +98,8 @@ def save_dataset_images(
Path(images_directory_path).mkdir(parents=True, exist_ok=True)
for image_path, image in images.items():
image_name = str(Path(image_path).stem)
image_ext = str(Path(image_path).suffix)
target_image_path = os.path.join(
images_directory_path, f"{image_name}{image_ext}"
)
image_name = Path(image_path).name
target_image_path = os.path.join(images_directory_path, image_name)
cv2.imwrite(target_image_path, image)