| import numpy as np | |
| import torch | |
| def get_tubes(masks_per_frame, tube_length): | |
| rp = torch.randperm(len(masks_per_frame)) | |
| masks_per_frame = masks_per_frame[rp] | |
| tubes = [masks_per_frame] | |
| for x in range(tube_length - 1): | |
| masks_per_frame = masks_per_frame.clone() | |
| rp = torch.randperm(len(masks_per_frame)) | |
| masks_per_frame = masks_per_frame[rp] | |
| tubes.append(masks_per_frame) | |
| tubes = torch.vstack(tubes) | |
| return tubes | |
| class RotatedTableMaskingGenerator: | |
| def __init__(self, | |
| input_size, | |
| mask_ratio, | |
| tube_length, | |
| batch_size, | |
| mask_type='rotated_table', | |
| seed=None, | |
| randomize_num_visible=False): | |
| self.batch_size = batch_size | |
| self.mask_ratio = mask_ratio | |
| self.tube_length = tube_length | |
| self.frames, self.height, self.width = input_size | |
| self.num_patches_per_frame = self.height * self.width | |
| self.total_patches = self.frames * self.num_patches_per_frame | |
| self.seed = seed | |
| self.randomize_num_visible = randomize_num_visible | |
| self.mask_type = mask_type | |
| def __repr__(self): | |
| repr_str = "Inverted Table Mask: total patches {}, tube length {}, randomize num visible? {}, seed {}".format( | |
| self.total_patches, self.tube_length, self.randomize_num_visible, self.seed | |
| ) | |
| return repr_str | |
| def __call__(self, m=None): | |
| if self.mask_type == 'rotated_table_magvit': | |
| self.mask_ratio = np.random.uniform(low=0.0, high=1) | |
| self.mask_ratio = np.cos(self.mask_ratio * np.pi / 2) | |
| elif self.mask_type == 'rotated_table_maskvit': | |
| self.mask_ratio = np.random.uniform(low=0.5, high=1) | |
| all_masks = [] | |
| for b in range(self.batch_size): | |
| self.num_masks_per_frame = max(0, int(self.mask_ratio * self.num_patches_per_frame)) | |
| self.total_masks = self.tube_length * self.num_masks_per_frame | |
| num_masks = self.num_masks_per_frame | |
| if self.randomize_num_visible: | |
| assert "Randomize num visible Not implemented" | |
| num_masks = self.rng.randint(low=num_masks, high=(self.num_patches_per_frame + 1)) | |
| if self.mask_ratio == 0: | |
| mask_per_frame = torch.hstack([ | |
| torch.zeros(self.num_patches_per_frame - num_masks), | |
| ]) | |
| else: | |
| mask_per_frame = torch.hstack([ | |
| torch.zeros(self.num_patches_per_frame - num_masks), | |
| torch.ones(num_masks), | |
| ]) | |
| tubes = get_tubes(mask_per_frame, self.tube_length) | |
| top = torch.zeros(self.height * self.width).to(tubes.dtype) | |
| top = torch.tile(top, (self.frames - self.tube_length, 1)) | |
| mask = torch.cat([top, tubes]) | |
| mask = mask.flatten() | |
| all_masks.append(mask) | |
| return torch.stack(all_masks) | |