Spaces:
Runtime error
Runtime error
# 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 | |
class Flowers102(BaseDataset): | |
"""The Oxford 102 Flower Dataset. | |
Support the `Oxford 102 Flowers Dataset <https://www.robots.ox.ac.uk/~vgg/data/flowers/102/>`_ Dataset. | |
After downloading and decompression, the dataset directory structure is as follows. | |
Flowers102 dataset directory: :: | |
Flowers102 | |
βββ jpg | |
β βββ image_00001.jpg | |
β βββ image_00002.jpg | |
β βββ ... | |
βββ imagelabels.mat | |
βββ setid.mat | |
βββ ... | |
Args: | |
data_root (str): The root directory for Oxford 102 Flowers dataset. | |
split (str, optional): The dataset split, supports "train", | |
"val", "trainval", and "test". Default to "trainval". | |
Examples: | |
>>> from mmpretrain.datasets import Flowers102 | |
>>> train_dataset = Flowers102(data_root='data/Flowers102', split='trainval') | |
>>> train_dataset | |
Dataset Flowers102 | |
Number of samples: 2040 | |
Root of dataset: data/Flowers102 | |
>>> test_dataset = Flowers102(data_root='data/Flowers102', split='test') | |
>>> test_dataset | |
Dataset Flowers102 | |
Number of samples: 6149 | |
Root of dataset: data/Flowers102 | |
""" # noqa: E501 | |
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 | |
ann_file = 'imagelabels.mat' | |
data_prefix = 'jpg' | |
train_test_split_file = 'setid.mat' | |
test_mode = split == 'test' | |
self.backend = get_file_backend(data_root, enable_singleton=True) | |
self.train_test_split_file = self.backend.join_path( | |
data_root, train_test_split_file) | |
super(Flowers102, 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.""" | |
label_dict = mat4py.loadmat(self.ann_file)['labels'] | |
split_list = mat4py.loadmat(self.train_test_split_file) | |
if self.split == 'train': | |
split_list = split_list['trnid'] | |
elif self.split == 'val': | |
split_list = split_list['valid'] | |
elif self.split == 'test': | |
split_list = split_list['tstid'] | |
else: | |
train_ids = split_list['trnid'] | |
val_ids = split_list['valid'] | |
train_ids.extend(val_ids) | |
split_list = train_ids | |
data_list = [] | |
for sample_id in split_list: | |
img_name = 'image_%05d.jpg' % (sample_id) | |
img_path = self.backend.join_path(self.img_prefix, img_name) | |
gt_label = int(label_dict[sample_id - 1]) - 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 | |