update from_pascal_voc to match v2 loader

This commit is contained in:
kirilllzaitsev committed 2023-07-26 10:21:17 +02:00
1 parent d7e08d42eb
commit 3723eae50c
1 file changed
+12 -26
+12 -26
View File
@@ -197,7 +197,10 @@ class DetectionDataset(BaseDataset):
@classmethod
def from_pascal_voc(
cls, images_directory_path: str, annotations_directory_path: str
cls,
images_directory_path: str,
annotations_directory_path: str,
force_masks: bool = False,
) -> DetectionDataset:
"""
Creates a Dataset instance from PASCAL VOC formatted data.
@@ -231,34 +234,17 @@ class DetectionDataset(BaseDataset):
['dog', 'person']
```
"""
image_paths = list_files_with_extensions(
directory=images_directory_path, extensions=["jpg", "jpeg", "png"]
)
annotation_paths = list_files_with_extensions(
directory=annotations_directory_path, extensions=["xml"]
classes, images, annotations = load_pascal_voc_annotations(
images_directory_path=images_directory_path,
annotations_directory_path=annotations_directory_path,
force_masks=force_masks,
)
raw_annotations: List[Tuple[str, Detections, List[str]]] = [
load_pascal_voc_annotations(annotation_path=str(annotation_path))
for annotation_path in annotation_paths
]
for annotation in annotations.values():
class_id = [classes.index(class_name) for class_name in annotation.class_id]
annotation.class_id = np.array(class_id)
classes = []
for annotation in raw_annotations:
classes.extend(annotation[2])
classes = list(set(classes))
for annotation in raw_annotations:
class_id = [classes.index(class_name) for class_name in annotation[2]]
annotation[1].class_id = np.array(class_id)
images = {
image_path.name: cv2.imread(str(image_path)) for image_path in image_paths
}
annotations = {
image_name: detections for image_name, detections, _ in raw_annotations
}
return DetectionDataset(classes=classes, images=images, annotations=annotations)
@classmethod