Skip to content

Commit c399c3f

Browse files
authored
add background category to prototype VOC dataset (#5687)
1 parent 9a9ae1e commit c399c3f

File tree

2 files changed

+6
-1
lines changed

2 files changed

+6
-1
lines changed

torchvision/prototype/datasets/_builtin/voc.categories

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
__background__
12
aeroplane
23
bicycle
34
bird

torchvision/prototype/datasets/_builtin/voc.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -219,4 +219,8 @@ def _generate_categories(self) -> List[str]:
219219
dp = Filter(archive_dp, self._filter_anns)
220220
dp = Mapper(dp, self._parse_detection_ann, input_col=1)
221221

222-
return sorted({instance["name"] for _, anns in dp for instance in anns["object"]})
222+
categories = sorted({instance["name"] for _, anns in dp for instance in anns["object"]})
223+
# We add a background category to be used during segmentation
224+
categories.insert(0, "__background__")
225+
226+
return categories

0 commit comments

Comments
 (0)