Spaces:
Sleeping
Sleeping
File size: 3,055 Bytes
b34d1d6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
from mmdet.registry import DATASETS
from mmdet.datasets.base_video_dataset import BaseVideoDataset
@DATASETS.register_module()
class YouTubeVISDatasetV2(BaseVideoDataset):
"""YouTube VIS dataset for video instance segmentation.
Args:
dataset_version (str): Select dataset year version.
"""
def __init__(self, dataset_version: str, *args, **kwargs):
self.set_dataset_classes(dataset_version)
self.dataset_name = f'YouTubeVISDataset_{dataset_version}'
super().__init__(*args, **kwargs)
@classmethod
def set_dataset_classes(cls, dataset_version: str) -> None:
"""Pass the category of the corresponding year to metainfo.
Args:
dataset_version (str): Select dataset year version.
"""
classes_2019_version = ('person', 'giant_panda', 'lizard', 'parrot',
'skateboard', 'sedan', 'ape', 'dog', 'snake',
'monkey', 'hand', 'rabbit', 'duck', 'cat',
'cow', 'fish', 'train', 'horse', 'turtle',
'bear', 'motorbike', 'giraffe', 'leopard',
'fox', 'deer', 'owl', 'surfboard', 'airplane',
'truck', 'zebra', 'tiger', 'elephant',
'snowboard', 'boat', 'shark', 'mouse', 'frog',
'eagle', 'earless_seal', 'tennis_racket')
classes_2021_version = ('airplane', 'bear', 'bird', 'boat', 'car',
'cat', 'cow', 'deer', 'dog', 'duck',
'earless_seal', 'elephant', 'fish',
'flying_disc', 'fox', 'frog', 'giant_panda',
'giraffe', 'horse', 'leopard', 'lizard',
'monkey', 'motorbike', 'mouse', 'parrot',
'person', 'rabbit', 'shark', 'skateboard',
'snake', 'snowboard', 'squirrel', 'surfboard',
'tennis_racket', 'tiger', 'train', 'truck',
'turtle', 'whale', 'zebra')
classes_ovis_version = ('Person', 'Bird', 'Cat', 'Dog', 'Horse',
'Sheep', 'Cow', 'Elephant', 'Bear', 'Zebra',
'Giraffe', 'Poultry', 'Giant_panda', 'Lizard', 'Parrot',
'Monkey', 'Rabbit', 'Tiger', 'Fish', 'Turtle',
'Bicycle', 'Motorcycle', 'Airplane', 'Boat', 'Vehical')
if dataset_version == '2019':
cls.METAINFO = dict(classes=classes_2019_version)
elif dataset_version == '2021':
cls.METAINFO = dict(classes=classes_2021_version)
elif dataset_version == 'ovis':
cls.METAINFO = dict(classes=classes_ovis_version)
else:
raise NotImplementedError('Not supported YouTubeVIS dataset'
f'version: {dataset_version}')
|