Class MetadataTransform

fourm/data/modality_transforms.py:843–1006

      

class MetadataTransform(AbstractTransform):

    def __init__(self, 
                 special_vmin: int = 0, 
                 special_vmax: int = 999, 
                 shuffle: bool = True, 
                 random_trunc: bool = False, 
                 return_chunks: bool = True, 
                 return_raw: bool = False,
                 image_dim_bin_size: int = 32,):
        """Metadata transform that takes in a metadata dictionary and converts 
        it into a string, or list of strings (for chunked span masking).
        Uses special tokens v1 to denote metadata types, and v0 for their values.

        Args:
            special_vmin: Minimum value for special tokens
            special_vmax: Maximum value for special tokens
            shuffle: Whether to shuffle the metadata order
            random_trunc: Whether to randomly truncate the returned metadata
            return_chunks: Whether to return a list of strings (for chunked span masking),
                or a single string with all metadata concatenated
            return_raw: Whether to return the raw metadata dictionary
        """
        self.special_vmin = special_vmin
        self.special_vmax = special_vmax
        self.shuffle = shuffle
        self.random_trunc = random_trunc
        self.return_chunks = return_chunks
        self.return_raw = return_raw
        self.image_dim_bin_size = image_dim_bin_size

        # Explicit map to make sure that additional entries do not change existing IDs
        # TODO: Make this work with other text tokenizers
        self.metadata_id_map = {
            'original_width': 'v1=0',
            'original_height': 'v1=1',
            'caption_n_chars': 'v1=2',
            'caption_n_words': 'v1=3',
            'caption_n_sentences': 'v1=4',
            'n_humans': 'v1=5',
            'n_sam_instances': 'v1=6',
            'n_coco_instances': 'v1=7',
            'coco_instance_diversity': 'v1=8',
            'colorfulness': 'v1=9',
            'brightness': 'v1=10',
            'contrast': 'v1=11',
            'saturation': 'v1=12',
            'entropy': 'v1=13',
            'walkability': 'v1=14',
            'objectness': 'v1=15',
            'semantic_diversity': 'v1=16',
            'geometric_complexity': 'v1=17',
            'occlusion_score': 'v1=18',
            'watermark_score': 'v1=19',
            'aesthetic_score': 'v1=20',
        }
        self.id_metadata_map = {v: k for k, v in self.metadata_id_map.items()}