XaviXva/Video-LLaVA

1import cv22import torch3from PIL import Image4from torch import nn5from torchvision import transforms6from transformers import ProcessorMixin, BatchEncoding7from transformers.image_processing_utils import BatchFeature8 9OPENAI_DATASET_MEAN = (0.48145466, 0.4578275, 0.40821073)10OPENAI_DATASET_STD = (0.26862954, 0.26130258, 0.27577711)11 12def make_list_of_images(x):13    if not isinstance(x, list):14        return [x]15    return x16 17def opencv_loader(path):18    return cv2.imread(path, cv2.IMREAD_UNCHANGED).astype('float32')19 20 21class DepthNorm(nn.Module):22    def __init__(23        self,24        max_depth=0,25        min_depth=0.01,26    ):27        super().__init__()28        self.max_depth = max_depth29        self.min_depth = min_depth30        self.scale = 1000.0  # nyuv2 abs.depth31 32    def forward(self, image):33        # image = np.array(image)34        depth_img = image / self.scale  # (H, W)   in meters35        depth_img = depth_img.clip(min=self.min_depth)36        if self.max_depth != 0:37            depth_img = depth_img.clip(max=self.max_depth)38            depth_img /= self.max_depth   #  0-139        else:40            depth_img /= depth_img.max()41        depth_img = torch.from_numpy(depth_img).unsqueeze(0).repeat(3, 1, 1)  # assume image42        return depth_img.to(torch.get_default_dtype())43 44def get_depth_transform(config):45    config = config.vision_config46    transform = transforms.Compose(47        [48            DepthNorm(max_depth=config.max_depth),49            transforms.Resize(224, interpolation=transforms.InterpolationMode.BICUBIC),50            transforms.CenterCrop(224),51            transforms.Normalize(OPENAI_DATASET_MEAN, OPENAI_DATASET_STD),  # assume image52            # transforms.Normalize((0.5, ), (0.5, ))  # 0-1 to norm distribution53            # transforms.Normalize((0.0418, ), (0.0295, ))  # sun rgb-d  imagebind54            # transforms.Normalize((0.02, ), (0.00295, ))  # nyuv255        ]56    )57    return transform58 59def load_and_transform_depth(depth_path, transform):60    depth = opencv_loader(depth_path)61    depth_outputs = transform(depth)62    return depth_outputs63 64class LanguageBindDepthProcessor(ProcessorMixin):65    attributes = []66    tokenizer_class = ("LanguageBindDepthTokenizer")67 68    def __init__(self, config, tokenizer=None, **kwargs):69        super().__init__(**kwargs)70        self.config = config71        self.transform = get_depth_transform(config)72        self.image_processor = load_and_transform_depth73        self.tokenizer = tokenizer74 75    def __call__(self, images=None, text=None, context_length=77, return_tensors=None, **kwargs):76        if text is None and images is None:77            raise ValueError("You have to specify either text or images. Both cannot be none.")78 79        if text is not None:80            encoding = self.tokenizer(text, max_length=context_length, padding='max_length',81                                      truncation=True, return_tensors=return_tensors, **kwargs)82 83        if images is not None:84            images = make_list_of_images(images)85            image_features = [self.image_processor(image, self.transform) for image in images]86            image_features = torch.stack(image_features)87 88        if text is not None and images is not None:89            encoding["pixel_values"] = image_features90            return encoding91        elif text is not None:92            return encoding93        else:94            return {"pixel_values": image_features}95 96    def batch_decode(self, skip_special_tokens=True, *args, **kwargs):97        """98        This method forwards all its arguments to CLIPTokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please99        refer to the docstring of this method for more information.100        """101        return self.tokenizer.batch_decode(*args, skip_special_tokens=skip_special_tokens, **kwargs)102 103    def decode(self, skip_special_tokens=True, *args, **kwargs):104        """105        This method forwards all its arguments to CLIPTokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to106        the docstring of this method for more information.107        """108        return self.tokenizer.decode(*args, skip_special_tokens=skip_special_tokens, **kwargs)109