move fixing of class_ids to load_pascal_voc_annotations

This commit is contained in:
kirilllzaitsev 2023-07-26 10:46:31 +02:00
parent a91118322b
commit 6dde70eb0c
2 changed files with 4 additions and 4 deletions

View File

@ -241,10 +241,6 @@ class DetectionDataset(BaseDataset):
force_masks=force_masks,
)
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)
return DetectionDataset(classes=classes, images=images, annotations=annotations)
@classmethod

View File

@ -203,6 +203,10 @@ def load_pascal_voc_annotations(
classes = list(set(classes))
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)
return classes, images, annotations