# Copyright (c) OpenMMLab. All rights reserved.
import os
import numpy as np
from PIL import Image
from easycv.datasets.registry import DATASOURCES
from easycv.file import io
[docs]@DATASOURCES.register_module()
class ClsSourceCUB(object):
"""The CUB-200-2011 Dataset.
Support the `CUB-200-2011 <http://www.vision.caltech.edu/visipedia/CUB-200-2011.html>`_ Dataset.
Comparing with the `CUB-200 <http://www.vision.caltech.edu/visipedia/CUB-200.html>`_ Dataset,
there are much more pictures in `CUB-200-2011`.
Args:
ann_file (str): the annotation file.
images.txt in CUB.
image_class_labels_file (str): the label file.
image_class_labels.txt in CUB.
train_test_split_file (str): the split file.
train_test_split_file.txt in CUB.
""" # noqa: E501
CLASSES = [
'Black_footed_Albatross', 'Laysan_Albatross', 'Sooty_Albatross',
'Groove_billed_Ani', 'Crested_Auklet', 'Least_Auklet',
'Parakeet_Auklet', 'Rhinoceros_Auklet', 'Brewer_Blackbird',
'Red_winged_Blackbird', 'Rusty_Blackbird', 'Yellow_headed_Blackbird',
'Bobolink', 'Indigo_Bunting', 'Lazuli_Bunting', 'Painted_Bunting',
'Cardinal', 'Spotted_Catbird', 'Gray_Catbird', 'Yellow_breasted_Chat',
'Eastern_Towhee', 'Chuck_will_Widow', 'Brandt_Cormorant',
'Red_faced_Cormorant', 'Pelagic_Cormorant', 'Bronzed_Cowbird',
'Shiny_Cowbird', 'Brown_Creeper', 'American_Crow', 'Fish_Crow',
'Black_billed_Cuckoo', 'Mangrove_Cuckoo', 'Yellow_billed_Cuckoo',
'Gray_crowned_Rosy_Finch', 'Purple_Finch', 'Northern_Flicker',
'Acadian_Flycatcher', 'Great_Crested_Flycatcher', 'Least_Flycatcher',
'Olive_sided_Flycatcher', 'Scissor_tailed_Flycatcher',
'Vermilion_Flycatcher', 'Yellow_bellied_Flycatcher', 'Frigatebird',
'Northern_Fulmar', 'Gadwall', 'American_Goldfinch',
'European_Goldfinch', 'Boat_tailed_Grackle', 'Eared_Grebe',
'Horned_Grebe', 'Pied_billed_Grebe', 'Western_Grebe', 'Blue_Grosbeak',
'Evening_Grosbeak', 'Pine_Grosbeak', 'Rose_breasted_Grosbeak',
'Pigeon_Guillemot', 'California_Gull', 'Glaucous_winged_Gull',
'Heermann_Gull', 'Herring_Gull', 'Ivory_Gull', 'Ring_billed_Gull',
'Slaty_backed_Gull', 'Western_Gull', 'Anna_Hummingbird',
'Ruby_throated_Hummingbird', 'Rufous_Hummingbird', 'Green_Violetear',
'Long_tailed_Jaeger', 'Pomarine_Jaeger', 'Blue_Jay', 'Florida_Jay',
'Green_Jay', 'Dark_eyed_Junco', 'Tropical_Kingbird', 'Gray_Kingbird',
'Belted_Kingfisher', 'Green_Kingfisher', 'Pied_Kingfisher',
'Ringed_Kingfisher', 'White_breasted_Kingfisher',
'Red_legged_Kittiwake', 'Horned_Lark', 'Pacific_Loon', 'Mallard',
'Western_Meadowlark', 'Hooded_Merganser', 'Red_breasted_Merganser',
'Mockingbird', 'Nighthawk', 'Clark_Nutcracker',
'White_breasted_Nuthatch', 'Baltimore_Oriole', 'Hooded_Oriole',
'Orchard_Oriole', 'Scott_Oriole', 'Ovenbird', 'Brown_Pelican',
'White_Pelican', 'Western_Wood_Pewee', 'Sayornis', 'American_Pipit',
'Whip_poor_Will', 'Horned_Puffin', 'Common_Raven',
'White_necked_Raven', 'American_Redstart', 'Geococcyx',
'Loggerhead_Shrike', 'Great_Grey_Shrike', 'Baird_Sparrow',
'Black_throated_Sparrow', 'Brewer_Sparrow', 'Chipping_Sparrow',
'Clay_colored_Sparrow', 'House_Sparrow', 'Field_Sparrow',
'Fox_Sparrow', 'Grasshopper_Sparrow', 'Harris_Sparrow',
'Henslow_Sparrow', 'Le_Conte_Sparrow', 'Lincoln_Sparrow',
'Nelson_Sharp_tailed_Sparrow', 'Savannah_Sparrow', 'Seaside_Sparrow',
'Song_Sparrow', 'Tree_Sparrow', 'Vesper_Sparrow',
'White_crowned_Sparrow', 'White_throated_Sparrow',
'Cape_Glossy_Starling', 'Bank_Swallow', 'Barn_Swallow',
'Cliff_Swallow', 'Tree_Swallow', 'Scarlet_Tanager', 'Summer_Tanager',
'Artic_Tern', 'Black_Tern', 'Caspian_Tern', 'Common_Tern',
'Elegant_Tern', 'Forsters_Tern', 'Least_Tern', 'Green_tailed_Towhee',
'Brown_Thrasher', 'Sage_Thrasher', 'Black_capped_Vireo',
'Blue_headed_Vireo', 'Philadelphia_Vireo', 'Red_eyed_Vireo',
'Warbling_Vireo', 'White_eyed_Vireo', 'Yellow_throated_Vireo',
'Bay_breasted_Warbler', 'Black_and_white_Warbler',
'Black_throated_Blue_Warbler', 'Blue_winged_Warbler', 'Canada_Warbler',
'Cape_May_Warbler', 'Cerulean_Warbler', 'Chestnut_sided_Warbler',
'Golden_winged_Warbler', 'Hooded_Warbler', 'Kentucky_Warbler',
'Magnolia_Warbler', 'Mourning_Warbler', 'Myrtle_Warbler',
'Nashville_Warbler', 'Orange_crowned_Warbler', 'Palm_Warbler',
'Pine_Warbler', 'Prairie_Warbler', 'Prothonotary_Warbler',
'Swainson_Warbler', 'Tennessee_Warbler', 'Wilson_Warbler',
'Worm_eating_Warbler', 'Yellow_Warbler', 'Northern_Waterthrush',
'Louisiana_Waterthrush', 'Bohemian_Waxwing', 'Cedar_Waxwing',
'American_Three_toed_Woodpecker', 'Pileated_Woodpecker',
'Red_bellied_Woodpecker', 'Red_cockaded_Woodpecker',
'Red_headed_Woodpecker', 'Downy_Woodpecker', 'Bewick_Wren',
'Cactus_Wren', 'Carolina_Wren', 'House_Wren', 'Marsh_Wren',
'Rock_Wren', 'Winter_Wren', 'Common_Yellowthroat'
]
[docs] def __init__(self, *args, ann_file, image_class_labels_file,
train_test_split_file, test_mode, data_prefix, **kwargs):
self.image_class_labels_file = image_class_labels_file
self.train_test_split_file = train_test_split_file
self.ann_file = ann_file
self.test_mode = test_mode
self.data_prefix = data_prefix
self.data_infos = self.load_annotations()
[docs] def load_annotations(self):
with open(self.ann_file) as f:
samples = [x.strip().split(' ')[1] for x in f.readlines()]
with open(self.image_class_labels_file) as f:
gt_labels = [
# in the official CUB-200-2011 dataset, labels in
# image_class_labels_file are started from 1, so
# here we need to '- 1' to let them start from 0.
int(x.strip().split(' ')[1]) - 1 for x in f.readlines()
]
with open(self.train_test_split_file) as f:
splits = [int(x.strip().split(' ')[1]) for x in f.readlines()]
assert len(samples) == len(gt_labels) == len(splits),\
f'samples({len(samples)}), gt_labels({len(gt_labels)}) and ' \
f'splits({len(splits)}) should have same length.'
data_infos = []
for filename, gt_label, split in zip(samples, gt_labels, splits):
if split and self.test_mode:
# skip train samples when test_mode=True
continue
elif not split and not self.test_mode:
# skip test samples when test_mode=False
continue
info = {}
info['img_info'] = os.path.join(self.data_prefix, filename)
info['gt_label'] = np.array(gt_label, dtype=np.int64)
data_infos.append(info)
return data_infos
def __len__(self):
return len(self.data_infos)
def __getitem__(self, idx):
data_info = self.data_infos[idx]
img = Image.open(io.open(data_info['img_info'], 'rb'))
if img.mode != 'RGB':
img = img.convert('RGB')
label = data_info['gt_label']
result_dict = {'img': img, 'gt_labels': label}
return result_dict