# -------------------------------------------------------- # Based on BEiT, timm, DINO and DeiT code bases # https://github.com/microsoft/unilm/tree/master/beit # https://github.com/rwightman/pytorch-image-models/tree/master/timm # https://github.com/facebookresearch/deit # https://github.com/facebookresearch/dino # --------------------------------------------------------' import os from torchvision import transforms from .datasets import RawFrameClsDataset, VideoClsDataset from .masking_generator import ( RunningCellMaskingGenerator, TubeMaskingGenerator, ) from .pretrain_datasets import HybridVideoMAE, VideoMAE # noqa: F401 from .transforms import ( GroupMultiScaleCrop, GroupNormalize, Stack, ToTorchFormatTensor, ) class DataAugmentationForVideoMAEv2(object): def __init__(self, args): self.input_mean = [0.485, 0.456, 0.406] self.input_std = [0.229, 0.224, 0.225] div = True roll = False normalize = GroupNormalize(self.input_mean, self.input_std) self.train_augmentation = GroupMultiScaleCrop(args.input_size, [1, .875, .75, .66]) self.transform = transforms.Compose([ self.train_augmentation, Stack(roll=roll), ToTorchFormatTensor(div=div), normalize, ]) if args.mask_type == 'tube': self.encoder_mask_map_generator = TubeMaskingGenerator( args.window_size, args.mask_ratio) else: raise NotImplementedError( 'Unsupported encoder masking strategy type.') if args.decoder_mask_ratio > 0.: if args.decoder_mask_type == 'run_cell': self.decoder_mask_map_generator = RunningCellMaskingGenerator( args.window_size, args.decoder_mask_ratio) else: raise NotImplementedError( 'Unsupported decoder masking strategy type.') def __call__(self, images): process_data, _ = self.transform(images) encoder_mask_map = self.encoder_mask_map_generator() if hasattr(self, 'decoder_mask_map_generator'): decoder_mask_map = self.decoder_mask_map_generator() else: decoder_mask_map = 1 - encoder_mask_map return process_data, encoder_mask_map, decoder_mask_map def __repr__(self): repr = "(DataAugmentationForVideoMAEv2,\n" repr += " transform = %s,\n" % str(self.transform) repr += " Encoder Masking Generator = %s,\n" % str( self.encoder_mask_map_generator) if hasattr(self, 'decoder_mask_map_generator'): repr += " Decoder Masking Generator = %s,\n" % str( self.decoder_mask_map_generator) else: repr += " Do not use decoder masking,\n" repr += ")" return repr def build_pretraining_dataset(args): transform = DataAugmentationForVideoMAEv2(args) dataset = VideoMAE( root=args.data_root, setting=args.video_data_path, train=True, test_mode=False, name_pattern=args.fname_tmpl, video_ext='mp4', is_color=True, modality='rgb', num_segments=1, num_crop=1, new_length=args.num_frames, new_step=args.sampling_rate, transform=transform, temporal_jitter=False, lazy_init=False, num_sample=args.num_sample) # print("Data Aug = %s" % str(transform)) return dataset def build_dataset(is_train, test_mode, args): if is_train: mode = 'train' anno_path = os.path.join(args.video_data_path, 'train.csv') elif test_mode: mode = 'test' anno_path = os.path.join(args.video_data_path, 'val.csv') else: mode = 'validation' anno_path = os.path.join(args.video_data_path, 'val.csv') if args.data_set == 'Kinetics-400': if not args.sparse_sample: dataset = VideoClsDataset( anno_path=anno_path, data_root=args.data_root, mode=mode, clip_len=args.num_frames, frame_sample_rate=args.sampling_rate, num_segment=1, test_num_segment=args.test_num_segment, test_num_crop=args.test_num_crop, num_crop=1 if not test_mode else 3, keep_aspect_ratio=True, crop_size=args.input_size, short_side_size=args.short_side_size, new_height=256, new_width=320, sparse_sample=False, args=args) else: dataset = VideoClsDataset( anno_path=anno_path, data_root=args.data_root, mode=mode, clip_len=1, frame_sample_rate=1, num_segment=args.num_frames, test_num_segment=args.test_num_segment, test_num_crop=args.test_num_crop, num_crop=1 if not test_mode else 3, keep_aspect_ratio=True, crop_size=args.input_size, short_side_size=args.short_side_size, new_height=256, new_width=320, sparse_sample=True, args=args) nb_classes = 400 elif args.data_set == 'Kinetics-600': dataset = VideoClsDataset( anno_path=anno_path, data_root=args.data_root, mode=mode, clip_len=args.num_frames, frame_sample_rate=args.sampling_rate, num_segment=1, test_num_segment=args.test_num_segment, test_num_crop=args.test_num_crop, num_crop=1 if not test_mode else 3, keep_aspect_ratio=True, crop_size=args.input_size, short_side_size=args.short_side_size, new_height=256, new_width=320, args=args) nb_classes = 600 elif args.data_set == 'Kinetics-700': dataset = VideoClsDataset( anno_path=anno_path, data_root=args.data_root, mode=mode, clip_len=args.num_frames, frame_sample_rate=args.sampling_rate, num_segment=1, test_num_segment=args.test_num_segment, test_num_crop=args.test_num_crop, num_crop=1 if not test_mode else 3, keep_aspect_ratio=True, crop_size=args.input_size, short_side_size=args.short_side_size, new_height=256, new_width=320, args=args) nb_classes = 700 elif args.data_set == 'Kinetics-710': dataset = VideoClsDataset( anno_path=anno_path, data_root=args.data_root, mode=mode, clip_len=args.num_frames, frame_sample_rate=args.sampling_rate, num_segment=1, test_num_segment=args.test_num_segment, test_num_crop=args.test_num_crop, num_crop=1 if not test_mode else 3, keep_aspect_ratio=True, crop_size=args.input_size, short_side_size=args.short_side_size, new_height=256, new_width=320, args=args) nb_classes = 710 elif args.data_set == 'SSV2': dataset = VideoClsDataset( anno_path=anno_path, data_root=args.data_root, mode=mode, clip_len=args.num_frames, frame_sample_rate=args.sampling_rate, num_segment=1, test_num_segment=args.test_num_segment, test_num_crop=args.test_num_crop, num_crop=1 if not test_mode else 3, keep_aspect_ratio=True, crop_size=args.input_size, short_side_size=args.short_side_size, new_height=256, new_width=320, args=args) nb_classes = 174 # elif args.data_set == 'SSV2': # dataset = RawFrameClsDataset( # anno_path=anno_path, # data_root=args.data_root, # mode=mode, # clip_len=1, # num_segment=args.num_frames, # test_num_segment=args.test_num_segment, # test_num_crop=args.test_num_crop, # num_crop=1 if not test_mode else 3, # keep_aspect_ratio=True, # crop_size=args.input_size, # short_side_size=args.short_side_size, # new_height=256, # new_width=320, # filename_tmpl=args.fname_tmpl, # start_idx=args.start_idx, # args=args) # # nb_classes = 174 elif args.data_set == 'UCF101': dataset = VideoClsDataset( anno_path=anno_path, data_root=args.data_root, mode=mode, clip_len=args.num_frames, frame_sample_rate=args.sampling_rate, num_segment=1, test_num_segment=args.test_num_segment, test_num_crop=args.test_num_crop, num_crop=1 if not test_mode else 3, keep_aspect_ratio=True, crop_size=args.input_size, short_side_size=args.short_side_size, new_height=256, new_width=320, args=args) nb_classes = 101 elif args.data_set == 'HMDB51': dataset = VideoClsDataset( anno_path=anno_path, data_root=args.data_root, mode=mode, clip_len=args.num_frames, frame_sample_rate=args.sampling_rate, num_segment=1, test_num_segment=args.test_num_segment, test_num_crop=args.test_num_crop, num_crop=1 if not test_mode else 3, keep_aspect_ratio=True, crop_size=args.input_size, short_side_size=args.short_side_size, new_height=256, new_width=320, args=args) nb_classes = 51 elif args.data_set == 'Diving48': dataset = VideoClsDataset( anno_path=anno_path, data_root=args.data_root, mode=mode, clip_len=args.num_frames, frame_sample_rate=args.sampling_rate, num_segment=1, test_num_segment=args.test_num_segment, test_num_crop=args.test_num_crop, num_crop=1 if not test_mode else 3, keep_aspect_ratio=True, crop_size=args.input_size, short_side_size=args.short_side_size, new_height=256, new_width=320, args=args) nb_classes = 48 elif args.data_set == 'MIT': if not args.sparse_sample: dataset = VideoClsDataset( anno_path=anno_path, data_root=args.data_root, mode=mode, clip_len=args.num_frames, frame_sample_rate=args.sampling_rate, num_segment=1, test_num_segment=args.test_num_segment, test_num_crop=args.test_num_crop, num_crop=1 if not test_mode else 3, keep_aspect_ratio=True, crop_size=args.input_size, short_side_size=args.short_side_size, new_height=256, new_width=320, sparse_sample=False, args=args) else: dataset = VideoClsDataset( anno_path=anno_path, data_root=args.data_root, mode=mode, clip_len=1, frame_sample_rate=1, num_segment=args.num_frames, test_num_segment=args.test_num_segment, test_num_crop=args.test_num_crop, num_crop=1 if not test_mode else 3, keep_aspect_ratio=True, crop_size=args.input_size, short_side_size=args.short_side_size, new_height=256, new_width=320, sparse_sample=True, args=args) nb_classes = 339 else: raise NotImplementedError('Unsupported Dataset') assert nb_classes == args.nb_classes # print("Number of the class = %d" % args.nb_classes) return dataset, nb_classes