update from_pascal_voc to match v2 loader
This commit is contained in:
1 parent
d7e08d42eb
commit
3723eae50c
1 file changed
+12
-26
+12
-26
@@ -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
|
||||
|
||||
Reference in new issue
Block a user