images, texts = show_test_images(args)

similarity = logits_per_text.cpu().numpy() / 100
_ = show_similarity(images, texts, similarity, args)

with open("ru-clip/clip/evaluate/classes.json", "r") as file:
    classes = json.load(file)
logits_per_image, logits_per_text = call_model(model, tokenizer, args, prepare_classes(classes), images)
show_topk_probs(images, classes, logits_per_image, args, k=5)

import subprocess

CUDA_version = [s for s in subprocess.check_output(["nvcc", "--version"]).decode("UTF-8").split(", ") if s.startswith("release")][0].split(" ")[-1]
print("CUDA version:", CUDA_version)

if CUDA_version == "10.0":
    torch_version_suffix = "+cu100"
elif CUDA_version == "10.1":
    torch_version_suffix = "+cu101"
elif CUDA_version == "10.2":
    torch_version_suffix = ""
else:
    torch_version_suffix = "+cu110"

!pip install torch==1.7.1{torch_version_suffix} torchvision==0.8.2{torch_version_suffix} -f https://download.pytorch.org/whl/torch_stable.html ftfy regex

import numpy as np
import torch

print("Torch version:", torch.__version__)
!pip install transformers==3.5.1
!pip install pymorphy2

!git clone https://github.com/sberbank-ai/ru-clip.git

%load_ext autoreload

%autoreload 2

import sys
sys.path.append("ru-clip/")

from clip.evaluate.utils import (
    get_text_batch, get_image_batch, get_tokenizer,
    show_test_images, show_similarity,
    prepare_classes, call_model,
    show_topk_probs,
    load_weights_only,
    get_topk_accuracy,
    show_topk_accuracy
)
from clip.evaluate.tvd import get_text_probs_from_dataset
import torch
import json

model, args = load_weights_only("ViT-B/32-small")
model = model.cuda().float().eval()

     predictions: {‘predictions’: ‘’, ... ‘res_image’: ‘’}

class KFServingRuClipModel(kfserving.KFModel):
    def __init__(self, name: str, model_path="ViT-B/32-small"):
        super().__init__(name)
        self.name = name
        self.ready = False
        self.model_path = model_path

    def load(self):
        self.model, self.args = load_weights_only(self.model_path)
        self.model = self.model.cuda().float().eval()
        self.tokenizer = get_tokenizer()
        self.ready = True
    
    def preprocess(self, ims):
        input_resolution = self.args.img_transform.transforms[0].size
        preprocess = Compose([
            Resize(input_resolution, interpolation=Image.BICUBIC),
            CenterCrop(input_resolution),
            ToTensor()
        ])
        images = []
        for img in ims:
            image = preprocess(img)
            images.append(image)
        return images
    
    def get_similarity_score(self, texts, images):
        """
        Find the most similar image to text.
        `texts`: array of texts or one text ["some_desc"]
        `images`: array of images.
        """
        input_ids, attention_mask = get_text_batch(["Это " + desc for desc in texts], self.tokenizer, self.args)
        images = torch.tensor(np.stack(images))
        img_input = images.to(images.device if self.args.cpu else torch.cuda.current_device())
    
        with torch.no_grad():
            logits_per_image, logits_per_text = self.model(
                img_input={"x": img_input},
                text_input={"x": input_ids, "attention_mask": attention_mask}
            )
        similarity = logits_per_text.cpu().numpy() / 100
        sim_plt = show_similarity(images, texts, similarity, self.args)
        result_img = create_image(sim_plt)
        return similarity, result_img
    
    def zero_shot_cls(self, images, txt_classes, topk):
        """
        Classify images using the cosine similarity (times 100) as the logits to the softmax operation.
        Show image - it returns the best classes
        """
        text_classes = prepare_classes(txt_classes)
        input_ids, attention_mask = get_text_batch(text_classes, self.tokenizer, self.args)
        images = torch.tensor(np.stack(images))
        img_input = images.to(images.device if self.args.cpu else torch.cuda.current_device())
        with torch.no_grad():
            logits_per_image, logits_per_text = self.model(
                img_input={"x": img_input},
                text_input={"x": input_ids, "attention_mask": attention_mask}
            )
        top_probs, plt = show_topk_probs(images, txt_classes, logits_per_image, self.args, k=topk)
        result_img = create_image(plt)
        return top_probs, result_img
    
    def predict(self, request: Dict) -> Dict:
        texts = request['instances'][0]["texts"]
        
        img_strs = request['instances'][0].get('images', None)
        if img_strs is not None:
            images = open_images_base64(img_strs)
        
        images_links = request['instances'][0].get('image_links', None)

        if images_links is not None:
            images = open_image_link(images_links)

        images_preproc = self.preprocess(images)
        mode = request['instances'][0].get('mode', 'similarity') # zero-shot
        
        error_msg = None
        if mode == 'similarity':
            preds, res_image = self.get_similarity_score(texts, images_preproc)
            predictions = { label: list(pred) for pred, label in zip(preds, texts) }
        elif mode == 'zero-shot':
            topk = request['instances'][0].get("k", 3)
            predictions, res_image = self.zero_shot_cls(images_preproc, texts, topk)
        else:
            predictions, error_msg, res_image = None, "Please, choose between two modes: `similarity`, `zero-shot`.", None
        if error_msg is None:
            return {"predictions": f"{predictions}", "mode": str(mode), "res_image": res_image, "error_message": error_msg}
        else:
            return {"predictions": f"{predictions}", "mode": str(mode), "res_image": res_image}

!pip install -r requirements.txt
from kfserving_ru_CLIP import KFServingRuClipModel

model = KFServingRuClipModel("kfserving-clip")
model.load()

# zero-shot and links

url_cat = "https://cs11.livemaster.ru/storage/topic/NxN/2c/9b/9cf0a41d13ecb11439e6145dff576315df83op.jpg?h=3KvOPndE06tlraLLSmkHPQ"
url_cat2 = "https://pbs.twimg.com/profile_images/560798448962633728/rDEdUfV_.jpeg"
url_dog = "https://ichef.bbci.co.uk/news/640/cpsprodpb/475B/production/_98776281_gettyimages-521697453.jpg"
result = model.predict({"instances": [{
        "texts": ["собачка", "кошка", "мышка", "машина", "стол", "дом", "жидкость"],
        "mode": "zero-shot",
        "k": 3,
        "image_links": [url_dog, url_cat, url_cat2]
    }]
})

"Res: ", result

## zero-shot and images in base64
result = model.predict({"instances": [{
        "texts": ["парень", "бокал вина", "девушка", "небо"],
        "mode": "zero-shot",
        "images": [man, vino, woman, nebo]
    }]
})

"Res: ", result

## similarity and links
result = model.predict({"instances": [{
        "texts": ["собачка", "кошка", "мышка"],
        "mode": "similarity",
        "image_links": [url_dog, url_cat, url_cat2]
    }]
})

"Res: ", result

## similarity and images in base64
result = model.predict({"instances": [{
        "texts": ["парень", "бокал вина", "девушка", "sky"],
        "mode": "similarity",
        "images": [man, vino, woman, nebo]
    }]
})

"Res: ", result

from io import BytesIO
import base64
from PIL import Image

Image.open(BytesIO(base64.b64decode(result["res_image"])))