|
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 |
|
|
