XaviXva/Video-LLaVA

1import torch2from PIL import Image3from torchvision import transforms4from transformers import ProcessorMixin, BatchEncoding5from transformers.image_processing_utils import BatchFeature6 7OPENAI_DATASET_MEAN = (0.48145466, 0.4578275, 0.40821073)8OPENAI_DATASET_STD = (0.26862954, 0.26130258, 0.27577711)9 10def make_list_of_images(x):11    if not isinstance(x, list):12        return [x]13    return x14 15def get_thermal_transform(config):16    config = config.vision_config17    transform = transforms.Compose(18        [19            transforms.ToTensor(),20            transforms.Resize(224, interpolation=transforms.InterpolationMode.BICUBIC),21            transforms.CenterCrop(224),22            transforms.Normalize(OPENAI_DATASET_MEAN, OPENAI_DATASET_STD)  # assume image23        ]24    )25    return transform26 27 28def load_and_transform_thermal(thermal_path, transform):29    thermal = Image.open(thermal_path)30    thermal_outputs = transform(thermal)31    return thermal_outputs32 33class LanguageBindThermalProcessor(ProcessorMixin):34    attributes = []35    tokenizer_class = ("LanguageBindThermalTokenizer")36 37    def __init__(self, config, tokenizer=None, **kwargs):38        super().__init__(**kwargs)39        self.config = config40        self.transform = get_thermal_transform(config)41        self.image_processor = load_and_transform_thermal42        self.tokenizer = tokenizer43 44    def __call__(self, images=None, text=None, context_length=77, return_tensors=None, **kwargs):45        if text is None and images is None:46            raise ValueError("You have to specify either text or images. Both cannot be none.")47 48        if text is not None:49            encoding = self.tokenizer(text, max_length=context_length, padding='max_length',50                                      truncation=True, return_tensors=return_tensors, **kwargs)51 52        if images is not None:53            images = make_list_of_images(images)54            image_features = [self.image_processor(image, self.transform) for image in images]55            image_features = torch.stack(image_features)56 57        if text is not None and images is not None:58            encoding["pixel_values"] = image_features59            return encoding60        elif text is not None:61            return encoding62        else:63            return {"pixel_values": image_features}64 65    def batch_decode(self, skip_special_tokens=True, *args, **kwargs):66        """67        This method forwards all its arguments to CLIPTokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please68        refer to the docstring of this method for more information.69        """70        return self.tokenizer.batch_decode(*args, skip_special_tokens=skip_special_tokens, **kwargs)71 72    def decode(self, skip_special_tokens=True, *args, **kwargs):73        """74        This method forwards all its arguments to CLIPTokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to75        the docstring of this method for more information.76        """77        return self.tokenizer.decode(*args, skip_special_tokens=skip_special_tokens, **kwargs)78