Source code for mmaction.datasets.pose_dataset

# Copyright (c) OpenMMLab. All rights reserved.
import os.path as osp
from typing import Callable, Dict, List, Optional, Union

import mmengine

from mmaction.registry import DATASETS
from .base import BaseActionDataset

[docs]@DATASETS.register_module() class PoseDataset(BaseActionDataset): """Pose dataset for action recognition. The dataset loads pose and apply specified transforms to return a dict containing pose information. The ann_file is a pickle file, the json file contains a list of annotations, the fields of an annotation include frame_dir(video_id), total_frames, label, kp, kpscore. Args: ann_file (str): Path to the annotation file. pipeline (list[dict | callable]): A sequence of data transforms. split (str, optional): The dataset split used. For UCF101 and HMDB51, allowed choices are 'train1', 'test1', 'train2', 'test2', 'train3', 'test3'. For NTURGB+D, allowed choices are 'xsub_train', 'xsub_val', 'xview_train', 'xview_val'. For NTURGB+D 120, allowed choices are 'xsub_train', 'xsub_val', 'xset_train', 'xset_val'. For FineGYM, allowed choices are 'train', 'val'. Defaults to None. """ def __init__(self, ann_file: str, pipeline: List[Union[Dict, Callable]], split: Optional[str] = None, **kwargs) -> None: self.split = split super().__init__( ann_file, pipeline=pipeline, modality='Pose', **kwargs)
[docs] def load_data_list(self) -> List[Dict]: """Load annotation file to get skeleton information.""" assert self.ann_file.endswith('.pkl') mmengine.exists(self.ann_file) data_list = mmengine.load(self.ann_file) if self.split is not None: split, annos = data_list['split'], data_list['annotations'] identifier = 'filename' if 'filename' in annos[0] else 'frame_dir' split = set(split[self.split]) data_list = [x for x in annos if x[identifier] in split] # Sometimes we may need to load video from the file if 'video' in self.data_prefix: for item in data_list: if 'filename' in item: item['filename'] = osp.join(self.data_prefix['video'], item['filename']) if 'frame_dir' in item: item['frame_dir'] = osp.join(self.data_prefix['video'], item['frame_dir']) return data_list
Read the Docs v: latest
On Read the Docs
Project Home

Free document hosting provided by Read the Docs.