Skip to content

GitLab

Unverified Commit e8b830fc authored by Francisco Massa's avatar Francisco Massa Committed by GitHub
Browse files

VOC2007 support test set (#1340)

parent 6de158c4
Show whitespace changes
Inline Side-by-side
Showing with 8 additions and 4 deletions
torchvision/datasets/voc.py
View file @ e8b830fc
......@@ -83,8 +83,10 @@ class VOCSegmentation(VisionDataset):
self.url = DATASET_YEAR_DICT[year]['url']
self.filename = DATASET_YEAR_DICT[year]['filename']
self.md5 = DATASET_YEAR_DICT[year]['md5']
self.image_set = verify_str_arg(image_set, "image_set",
("train", "trainval", "val"))
valid_sets = ["train", "trainval", "val"]
if year == "2007":
valid_sets.append("test")
self.image_set = verify_str_arg(image_set, "image_set", valid_sets)
base_dir = DATASET_YEAR_DICT[year]['base_dir']
voc_root = os.path.join(self.root, base_dir)
image_dir = os.path.join(voc_root, 'JPEGImages')
......@@ -160,8 +162,10 @@ class VOCDetection(VisionDataset):
self.url = DATASET_YEAR_DICT[year]['url']
self.filename = DATASET_YEAR_DICT[year]['filename']
self.md5 = DATASET_YEAR_DICT[year]['md5']
self.image_set = verify_str_arg(image_set, "image_set",
("train", "trainval", "val"))
valid_sets = ["train", "trainval", "val"]
if year == "2007":
valid_sets.append("test")
self.image_set = verify_str_arg(image_set, "image_set", valid_sets)
base_dir = DATASET_YEAR_DICT[year]['base_dir']
voc_root = os.path.join(self.root, base_dir)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment