117 lines
3.6 KiB
Python
117 lines
3.6 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
from typing import List
|
|
|
|
import mat4py
|
|
from mmengine import get_file_backend
|
|
|
|
from mmpretrain.registry import DATASETS
|
|
from .base_dataset import BaseDataset
|
|
from .categories import DTD_CATEGORIES
|
|
|
|
|
|
@DATASETS.register_module()
|
|
class DTD(BaseDataset):
|
|
"""The Describable Texture Dataset (DTD).
|
|
|
|
Support the `Describable Texture Dataset <https://www.robots.ox.ac.uk/~vgg/data/dtd/>`_ Dataset.
|
|
After downloading and decompression, the dataset directory structure is as follows.
|
|
|
|
DTD dataset directory: ::
|
|
|
|
dtd
|
|
├── images
|
|
│ ├── banded
|
|
| | ├──banded_0002.jpg
|
|
| | ├──banded_0004.jpg
|
|
| | └── ...
|
|
│ └── ...
|
|
├── imdb
|
|
│ └── imdb.mat
|
|
├── labels
|
|
| | ├──labels_joint_anno.txt
|
|
| | ├──test1.txt
|
|
| | ├──test2.txt
|
|
| | └── ...
|
|
│ └── ...
|
|
└── ....
|
|
|
|
Args:
|
|
data_root (str): The root directory for Describable Texture dataset.
|
|
split (str, optional): The dataset split, supports "train",
|
|
"val", "trainval", and "test". Default to "trainval".
|
|
|
|
Examples:
|
|
>>> from mmpretrain.datasets import DTD
|
|
>>> train_dataset = DTD(data_root='data/dtd', split='trainval')
|
|
>>> train_dataset
|
|
Dataset DTD
|
|
Number of samples: 3760
|
|
Number of categories: 47
|
|
Root of dataset: data/dtd
|
|
>>> test_dataset = DTD(data_root='data/dtd', split='test')
|
|
>>> test_dataset
|
|
Dataset DTD
|
|
Number of samples: 1880
|
|
Number of categories: 47
|
|
Root of dataset: data/dtd
|
|
""" # noqa: E501
|
|
|
|
METAINFO = {'classes': DTD_CATEGORIES}
|
|
|
|
def __init__(self, data_root: str, split: str = 'trainval', **kwargs):
|
|
|
|
splits = ['train', 'val', 'trainval', 'test']
|
|
assert split in splits, \
|
|
f"The split must be one of {splits}, but get '{split}'"
|
|
self.split = split
|
|
|
|
data_prefix = 'images'
|
|
test_mode = split == 'test'
|
|
|
|
self.backend = get_file_backend(data_root, enable_singleton=True)
|
|
ann_file = self.backend.join_path('imdb', 'imdb.mat')
|
|
|
|
super(DTD, self).__init__(
|
|
ann_file=ann_file,
|
|
data_root=data_root,
|
|
data_prefix=data_prefix,
|
|
test_mode=test_mode,
|
|
**kwargs)
|
|
|
|
def load_data_list(self):
|
|
"""Load images and ground truth labels."""
|
|
|
|
data = mat4py.loadmat(self.ann_file)['images']
|
|
names = data['name']
|
|
labels = data['class']
|
|
parts = data['set']
|
|
num = len(names)
|
|
assert num == len(labels) == len(parts), 'get error ann file'
|
|
|
|
if self.split == 'train':
|
|
target_set = {1}
|
|
elif self.split == 'val':
|
|
target_set = {2}
|
|
elif self.split == 'test':
|
|
target_set = {3}
|
|
else:
|
|
target_set = {1, 2}
|
|
|
|
data_list = []
|
|
for i in range(num):
|
|
if parts[i] in target_set:
|
|
img_name = names[i]
|
|
img_path = self.backend.join_path(self.img_prefix, img_name)
|
|
gt_label = labels[i] - 1
|
|
info = dict(img_path=img_path, gt_label=gt_label)
|
|
data_list.append(info)
|
|
|
|
return data_list
|
|
|
|
def extra_repr(self) -> List[str]:
|
|
"""The extra repr information of the dataset."""
|
|
body = [
|
|
f'Root of dataset: \t{self.data_root}',
|
|
]
|
|
return body
|