baseplate/vit-gpt2-image-captioning

1from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer2import torch3from PIL import Image4from typing import Dict, List, Any5import requests6 7 8class EndpointHandler():9    def __init__(self, path=""):10        model = VisionEncoderDecoderModel.from_pretrained(11            "nlpconnect/vit-gpt2-image-captioning")12        feature_extractor = ViTImageProcessor.from_pretrained(13            "nlpconnect/vit-gpt2-image-captioning")14        tokenizer = AutoTokenizer.from_pretrained(15            "nlpconnect/vit-gpt2-image-captioning")16 17        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")18        model.to(device)19        self.model = model20        self.feature_extractor = feature_extractor21        self.tokenizer = tokenizer22 23    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:24        """25       data args:26            inputs (:obj: `str`)27            date (:obj: `str`)28      Return:29            A :obj:`list` | `dict`: will be serialized and returned30        """31        # get inputs32        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")33        max_length = 12834        num_beams = 435        gen_kwargs = {"max_length": max_length, "num_beams": num_beams}36        image_paths = data.pop("image_paths", data)37        images = []38        for image_path in image_paths:39            response = requests.get(image_path)40            response.raise_for_status()  # Raise an exception if the request failed41 42            with open("temp", "wb") as f:43                f.write(response.content)44            i_image = Image.open("temp")45            if i_image.mode != "RGB":46                i_image = i_image.convert(mode="RGB")47 48            images.append(i_image)49 50        pixel_values = self.feature_extractor(51            images=images, return_tensors="pt").pixel_values52        pixel_values = pixel_values.to(device)53 54        output_ids = self.model.generate(pixel_values, **gen_kwargs)55 56        preds = self.tokenizer.batch_decode(57            output_ids, skip_special_tokens=True)58        preds = [pred.strip() for pred in preds]59        return preds60