| | import os |
| | import tarfile |
| |
|
| | import h5py |
| | import numpy as np |
| | import torch |
| | import torch.utils.data as data |
| | from PIL import Image |
| | from scipy import io |
| | from torchvision.datasets.utils import download_url |
| |
|
| | DATASET_YEAR_DICT = { |
| | "2012": { |
| | "url": "http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar", |
| | "filename": "VOCtrainval_11-May-2012.tar", |
| | "md5": "6cd6e144f989b92b3379bac3b3de84fd", |
| | "base_dir": "VOCdevkit/VOC2012", |
| | }, |
| | "2011": { |
| | "url": "http://host.robots.ox.ac.uk/pascal/VOC/voc2011/VOCtrainval_25-May-2011.tar", |
| | "filename": "VOCtrainval_25-May-2011.tar", |
| | "md5": "6c3384ef61512963050cb5d687e5bf1e", |
| | "base_dir": "TrainVal/VOCdevkit/VOC2011", |
| | }, |
| | "2010": { |
| | "url": "http://host.robots.ox.ac.uk/pascal/VOC/voc2010/VOCtrainval_03-May-2010.tar", |
| | "filename": "VOCtrainval_03-May-2010.tar", |
| | "md5": "da459979d0c395079b5c75ee67908abb", |
| | "base_dir": "VOCdevkit/VOC2010", |
| | }, |
| | "2009": { |
| | "url": "http://host.robots.ox.ac.uk/pascal/VOC/voc2009/VOCtrainval_11-May-2009.tar", |
| | "filename": "VOCtrainval_11-May-2009.tar", |
| | "md5": "59065e4b188729180974ef6572f6a212", |
| | "base_dir": "VOCdevkit/VOC2009", |
| | }, |
| | "2008": { |
| | "url": "http://host.robots.ox.ac.uk/pascal/VOC/voc2008/VOCtrainval_14-Jul-2008.tar", |
| | "filename": "VOCtrainval_11-May-2012.tar", |
| | "md5": "2629fa636546599198acfcfbfcf1904a", |
| | "base_dir": "VOCdevkit/VOC2008", |
| | }, |
| | "2007": { |
| | "url": "http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtrainval_06-Nov-2007.tar", |
| | "filename": "VOCtrainval_06-Nov-2007.tar", |
| | "md5": "c52e279531787c972589f7e41ab4ae64", |
| | "base_dir": "VOCdevkit/VOC2007", |
| | }, |
| | } |
| |
|
| |
|
| | class VOCSegmentation(data.Dataset): |
| | """`Pascal VOC <http://host.robots.ox.ac.uk/pascal/VOC/>`_ Segmentation Dataset. |
| | |
| | Args: |
| | root (string): Root directory of the VOC Dataset. |
| | year (string, optional): The dataset year, supports years 2007 to 2012. |
| | image_set (string, optional): Select the image_set to use, ``train``, ``trainval`` or ``val`` |
| | download (bool, optional): If true, downloads the dataset from the internet and |
| | puts it in root directory. If dataset is already downloaded, it is not |
| | downloaded again. |
| | transform (callable, optional): A function/transform that takes in an PIL image |
| | and returns a transformed version. E.g, ``transforms.RandomCrop`` |
| | target_transform (callable, optional): A function/transform that takes in the |
| | target and transforms it. |
| | """ |
| |
|
| | CLASSES = 20 |
| | CLASSES_NAMES = [ |
| | "aeroplane", |
| | "bicycle", |
| | "bird", |
| | "boat", |
| | "bottle", |
| | "bus", |
| | "car", |
| | "cat", |
| | "chair", |
| | "cow", |
| | "diningtable", |
| | "dog", |
| | "horse", |
| | "motorbike", |
| | "person", |
| | "potted-plant", |
| | "sheep", |
| | "sofa", |
| | "train", |
| | "tvmonitor", |
| | "ambigious", |
| | ] |
| |
|
| | def __init__( |
| | self, |
| | root, |
| | year="2012", |
| | image_set="train", |
| | download=False, |
| | transform=None, |
| | target_transform=None, |
| | ): |
| | self.root = os.path.expanduser(root) |
| | self.year = year |
| | self.url = DATASET_YEAR_DICT[year]["url"] |
| | self.filename = DATASET_YEAR_DICT[year]["filename"] |
| | self.md5 = DATASET_YEAR_DICT[year]["md5"] |
| | self.transform = transform |
| | self.target_transform = target_transform |
| | self.image_set = image_set |
| | 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") |
| | mask_dir = os.path.join(voc_root, "SegmentationClass") |
| |
|
| | if download: |
| | download_extract(self.url, self.root, self.filename, self.md5) |
| |
|
| | if not os.path.isdir(voc_root): |
| | raise RuntimeError( |
| | "Dataset not found or corrupted." |
| | + " You can use download=True to download it" |
| | ) |
| |
|
| | splits_dir = os.path.join(voc_root, "ImageSets/Segmentation") |
| |
|
| | split_f = os.path.join(splits_dir, image_set.rstrip("\n") + ".txt") |
| |
|
| | if not os.path.exists(split_f): |
| | raise ValueError( |
| | 'Wrong image_set entered! Please use image_set="train" ' |
| | 'or image_set="trainval" or image_set="val"' |
| | ) |
| |
|
| | with open(os.path.join(split_f), "r") as f: |
| | file_names = [x.strip() for x in f.readlines()] |
| |
|
| | self.images = [os.path.join(image_dir, x + ".jpg") for x in file_names] |
| | self.masks = [os.path.join(mask_dir, x + ".png") for x in file_names] |
| | assert len(self.images) == len(self.masks) |
| |
|
| | def __getitem__(self, index): |
| | """ |
| | Args: |
| | index (int): Index |
| | |
| | Returns: |
| | tuple: (image, target) where target is the image segmentation. |
| | """ |
| | img = Image.open(self.images[index]).convert("RGB") |
| | target = Image.open(self.masks[index]) |
| |
|
| | if self.transform is not None: |
| | img = self.transform(img) |
| |
|
| | if self.target_transform is not None: |
| | target = np.array(self.target_transform(target)).astype("int32") |
| | target[target == 255] = -1 |
| | target = torch.from_numpy(target).long() |
| |
|
| | return img, target |
| |
|
| | @staticmethod |
| | def _mask_transform(mask): |
| | target = np.array(mask).astype("int32") |
| | target[target == 255] = -1 |
| | return torch.from_numpy(target).long() |
| |
|
| | def __len__(self): |
| | return len(self.images) |
| |
|
| | @property |
| | def pred_offset(self): |
| | return 0 |
| |
|
| |
|
| | class VOCClassification(data.Dataset): |
| | """`Pascal VOC <http://host.robots.ox.ac.uk/pascal/VOC/>`_ Segmentation Dataset. |
| | |
| | Args: |
| | root (string): Root directory of the VOC Dataset. |
| | year (string, optional): The dataset year, supports years 2007 to 2012. |
| | image_set (string, optional): Select the image_set to use, ``train``, ``trainval`` or ``val`` |
| | download (bool, optional): If true, downloads the dataset from the internet and |
| | puts it in root directory. If dataset is already downloaded, it is not |
| | downloaded again. |
| | transform (callable, optional): A function/transform that takes in an PIL image |
| | and returns a transformed version. E.g, ``transforms.RandomCrop`` |
| | """ |
| |
|
| | CLASSES = 20 |
| |
|
| | def __init__( |
| | self, root, year="2012", image_set="train", download=False, transform=None |
| | ): |
| | self.root = os.path.expanduser(root) |
| | self.year = year |
| | self.url = DATASET_YEAR_DICT[year]["url"] |
| | self.filename = DATASET_YEAR_DICT[year]["filename"] |
| | self.md5 = DATASET_YEAR_DICT[year]["md5"] |
| | self.transform = transform |
| | self.image_set = image_set |
| | 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") |
| | mask_dir = os.path.join(voc_root, "SegmentationClass") |
| |
|
| | if download: |
| | download_extract(self.url, self.root, self.filename, self.md5) |
| |
|
| | if not os.path.isdir(voc_root): |
| | raise RuntimeError( |
| | "Dataset not found or corrupted." |
| | + " You can use download=True to download it" |
| | ) |
| |
|
| | splits_dir = os.path.join(voc_root, "ImageSets/Segmentation") |
| |
|
| | split_f = os.path.join(splits_dir, image_set.rstrip("\n") + ".txt") |
| |
|
| | if not os.path.exists(split_f): |
| | raise ValueError( |
| | 'Wrong image_set entered! Please use image_set="train" ' |
| | 'or image_set="trainval" or image_set="val"' |
| | ) |
| |
|
| | with open(os.path.join(split_f), "r") as f: |
| | file_names = [x.strip() for x in f.readlines()] |
| |
|
| | self.images = [os.path.join(image_dir, x + ".jpg") for x in file_names] |
| | self.masks = [os.path.join(mask_dir, x + ".png") for x in file_names] |
| | assert len(self.images) == len(self.masks) |
| |
|
| | def __getitem__(self, index): |
| | """ |
| | Args: |
| | index (int): Index |
| | |
| | Returns: |
| | tuple: (image, target) where target is the image segmentation. |
| | """ |
| | img = Image.open(self.images[index]).convert("RGB") |
| | target = Image.open(self.masks[index]) |
| |
|
| | |
| | |
| | if self.transform is not None: |
| | img, target = self.transform(img, target) |
| |
|
| | visible_classes = np.unique(target) |
| | labels = torch.zeros(self.CLASSES) |
| | for id in visible_classes: |
| | if id not in (0, 255): |
| | labels[id - 1].fill_(1) |
| |
|
| | return img, labels |
| |
|
| | def __len__(self): |
| | return len(self.images) |
| |
|
| |
|
| | class VOCSBDClassification(data.Dataset): |
| | """`Pascal VOC <http://host.robots.ox.ac.uk/pascal/VOC/>`_ Segmentation Dataset. |
| | |
| | Args: |
| | root (string): Root directory of the VOC Dataset. |
| | year (string, optional): The dataset year, supports years 2007 to 2012. |
| | image_set (string, optional): Select the image_set to use, ``train``, ``trainval`` or ``val`` |
| | download (bool, optional): If true, downloads the dataset from the internet and |
| | puts it in root directory. If dataset is already downloaded, it is not |
| | downloaded again. |
| | transform (callable, optional): A function/transform that takes in an PIL image |
| | and returns a transformed version. E.g, ``transforms.RandomCrop`` |
| | """ |
| |
|
| | CLASSES = 20 |
| |
|
| | def __init__( |
| | self, |
| | root, |
| | sbd_root, |
| | year="2012", |
| | image_set="train", |
| | download=False, |
| | transform=None, |
| | ): |
| | self.root = os.path.expanduser(root) |
| | self.sbd_root = os.path.expanduser(sbd_root) |
| | self.year = year |
| | self.url = DATASET_YEAR_DICT[year]["url"] |
| | self.filename = DATASET_YEAR_DICT[year]["filename"] |
| | self.md5 = DATASET_YEAR_DICT[year]["md5"] |
| | self.transform = transform |
| | self.image_set = image_set |
| | 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") |
| | mask_dir = os.path.join(voc_root, "SegmentationClass") |
| | sbd_image_dir = os.path.join(sbd_root, "img") |
| | sbd_mask_dir = os.path.join(sbd_root, "cls") |
| |
|
| | if download: |
| | download_extract(self.url, self.root, self.filename, self.md5) |
| |
|
| | if not os.path.isdir(voc_root): |
| | raise RuntimeError( |
| | "Dataset not found or corrupted." |
| | + " You can use download=True to download it" |
| | ) |
| |
|
| | splits_dir = os.path.join(voc_root, "ImageSets/Segmentation") |
| |
|
| | split_f = os.path.join(splits_dir, image_set.rstrip("\n") + ".txt") |
| | sbd_split = os.path.join(sbd_root, "train.txt") |
| |
|
| | if not os.path.exists(split_f): |
| | raise ValueError( |
| | 'Wrong image_set entered! Please use image_set="train" ' |
| | 'or image_set="trainval" or image_set="val"' |
| | ) |
| |
|
| | with open(os.path.join(split_f), "r") as f: |
| | voc_file_names = [x.strip() for x in f.readlines()] |
| |
|
| | with open(os.path.join(sbd_split), "r") as f: |
| | sbd_file_names = [x.strip() for x in f.readlines()] |
| |
|
| | self.images = [os.path.join(image_dir, x + ".jpg") for x in voc_file_names] |
| | self.images += [os.path.join(sbd_image_dir, x + ".jpg") for x in sbd_file_names] |
| | self.masks = [os.path.join(mask_dir, x + ".png") for x in voc_file_names] |
| | self.masks += [os.path.join(sbd_mask_dir, x + ".mat") for x in sbd_file_names] |
| | assert len(self.images) == len(self.masks) |
| |
|
| | def __getitem__(self, index): |
| | """ |
| | Args: |
| | index (int): Index |
| | |
| | Returns: |
| | tuple: (image, target) where target is the image segmentation. |
| | """ |
| | img = Image.open(self.images[index]).convert("RGB") |
| | mask_path = self.masks[index] |
| | if mask_path[-3:] == "mat": |
| | target = io.loadmat(mask_path, struct_as_record=False, squeeze_me=True)[ |
| | "GTcls" |
| | ].Segmentation |
| | target = Image.fromarray(target, mode="P") |
| | else: |
| | target = Image.open(self.masks[index]) |
| |
|
| | if self.transform is not None: |
| | img, target = self.transform(img, target) |
| |
|
| | visible_classes = np.unique(target) |
| | labels = torch.zeros(self.CLASSES) |
| | for id in visible_classes: |
| | if id not in (0, 255): |
| | labels[id - 1].fill_(1) |
| |
|
| | return img, labels |
| |
|
| | def __len__(self): |
| | return len(self.images) |
| |
|
| |
|
| | def download_extract(url, root, filename, md5): |
| | download_url(url, root, filename, md5) |
| | with tarfile.open(os.path.join(root, filename), "r") as tar: |
| | tar.extractall(path=root) |
| |
|
| |
|
| | class VOCResults(data.Dataset): |
| | CLASSES = 20 |
| | CLASSES_NAMES = [ |
| | "aeroplane", |
| | "bicycle", |
| | "bird", |
| | "boat", |
| | "bottle", |
| | "bus", |
| | "car", |
| | "cat", |
| | "chair", |
| | "cow", |
| | "diningtable", |
| | "dog", |
| | "horse", |
| | "motorbike", |
| | "person", |
| | "potted-plant", |
| | "sheep", |
| | "sofa", |
| | "train", |
| | "tvmonitor", |
| | "ambigious", |
| | ] |
| |
|
| | def __init__(self, path): |
| | super(VOCResults, self).__init__() |
| |
|
| | self.path = os.path.join(path, "results.hdf5") |
| | self.data = None |
| |
|
| | print("Reading dataset length...") |
| | with h5py.File(self.path, "r") as f: |
| | self.data_length = len(f["/image"]) |
| |
|
| | def __len__(self): |
| | return self.data_length |
| |
|
| | def __getitem__(self, item): |
| | if self.data is None: |
| | self.data = h5py.File(self.path, "r") |
| |
|
| | image = torch.tensor(self.data["image"][item]) |
| | vis = torch.tensor(self.data["vis"][item]) |
| | target = torch.tensor(self.data["target"][item]) |
| | class_pred = torch.tensor(self.data["class_pred"][item]) |
| |
|
| | return image, vis, target, class_pred |
| |
|
