Skip to content

GitLab

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

Change SBDataset to make it support already-downloaded datasets (#865)

parent ccbb3221
Show whitespace changes
Inline Side-by-side
Showing with 6 additions and 1 deletion
torchvision/datasets/sbd.py
View file @ 3b59c6ee
import os import os
import shutil
from .vision import VisionDataset from .vision import VisionDataset
import numpy as np import numpy as np
...@@ -72,12 +73,16 @@ class SBDataset(VisionDataset): ...@@ -72,12 +73,16 @@ class SBDataset(VisionDataset):
self.mode = mode self.mode = mode
self.num_classes = 20 self.num_classes = 20
sbd_root = os.path.join(self.root, "benchmark_RELEASE", "dataset") sbd_root = self.root
image_dir = os.path.join(sbd_root, 'img') image_dir = os.path.join(sbd_root, 'img')
mask_dir = os.path.join(sbd_root, 'cls') mask_dir = os.path.join(sbd_root, 'cls')
if download: if download:
download_extract(self.url, self.root, self.filename, self.md5) download_extract(self.url, self.root, self.filename, self.md5)
extracted_ds_root = os.path.join(self.root, "benchmark_RELEASE", "dataset")
for f in ["cls", "img", "inst", "train.txt", "val.txt"]:
old_path = os.path.join(extracted_ds_root, f)
shutil.move(old_path, sbd_root)
download_url(self.voc_train_url, sbd_root, self.voc_split_filename, download_url(self.voc_train_url, sbd_root, self.voc_split_filename,
self.voc_split_md5) self.voc_split_md5)
......
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