TTP / mmpretrain /datasets /flowers102.py
KyanChen's picture
Upload 1861 files
3b96cb1
# 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
@DATASETS.register_module()
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