129 lines
5.1 KiB
Python
129 lines
5.1 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# Copyright 2021-2022 The Alibaba Fundamental Vision Team Authors. All rights reserved.
|
|
from typing import Any, Dict
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from PIL import Image
|
|
from tokenizers import BertWordPieceTokenizer
|
|
from torchvision.transforms import Compose, Normalize, Resize, ToTensor
|
|
|
|
from modelscope.metainfo import Models
|
|
from modelscope.models.base import TorchModel
|
|
from modelscope.models.builder import MODELS
|
|
from modelscope.outputs import OutputKeys
|
|
from modelscope.preprocessors import LoadImage
|
|
from modelscope.utils.constant import ModelFile, Tasks
|
|
from modelscope.utils.logger import get_logger
|
|
from .utils import TEAM, BertWrapper, CLIPVisionWrapper, CrossLayer
|
|
|
|
logger = get_logger()
|
|
|
|
__all__ = ['TEAMForMultiModalSimilarity']
|
|
|
|
|
|
@MODELS.register_module(Tasks.multi_modal_similarity, module_name=Models.team)
|
|
class TEAMForMultiModalSimilarity(TorchModel):
|
|
|
|
def __init__(self, model_dir, device_id=0, *args, **kwargs):
|
|
super().__init__(
|
|
model_dir=model_dir, device_id=device_id, *args, **kwargs)
|
|
|
|
text_model = BertWrapper(
|
|
config_json='{}/text_config.json'.format(model_dir),
|
|
feat_dim=768,
|
|
token_dim=1024)
|
|
text_model.bert.cls = None
|
|
image_model = CLIPVisionWrapper()
|
|
|
|
self.model = TEAM(
|
|
text_model,
|
|
image_model,
|
|
pretrained='{}/{}'.format(model_dir,
|
|
ModelFile.TORCH_MODEL_BIN_FILE))
|
|
self.model.eval()
|
|
|
|
self.device_id = device_id
|
|
if self.device_id >= 0 and torch.cuda.is_available():
|
|
self.model.to('cuda:{}'.format(self.device_id))
|
|
logger.info('Use GPU: {}'.format(self.device_id))
|
|
else:
|
|
self.device_id = -1
|
|
logger.info('Use CPU for inference')
|
|
|
|
self.text_tokenizer = BertWordPieceTokenizer(
|
|
'{}/{}'.format(model_dir, ModelFile.VOCAB_FILE), lowercase=False)
|
|
self.text_tokenizer.enable_truncation(max_length=30)
|
|
|
|
norm_op = Normalize((0.48145466, 0.4578275, 0.40821073),
|
|
(0.26862954, 0.26130258, 0.27577711))
|
|
self.img_preprocessor = Compose([
|
|
Resize((224, 224), interpolation=Image.BICUBIC),
|
|
ToTensor(), norm_op
|
|
])
|
|
|
|
def tokenize_text(self, text_str):
|
|
tokens = self.text_tokenizer.encode(text_str)
|
|
max_tokens = 30
|
|
text_ids_tensor = torch.zeros((1, max_tokens)).long()
|
|
text_mask_tensor = torch.zeros((1, max_tokens))
|
|
text_ids, text_mask = tokens.ids, tokens.attention_mask
|
|
text_ids_tensor[0, 0:len(text_ids)] = torch.tensor(text_ids)
|
|
text_mask_tensor[0, 0:len(text_mask)] = torch.tensor(text_mask)
|
|
return text_ids_tensor, text_mask_tensor
|
|
|
|
def forward(self, input: Dict[str, Any]) -> Dict[str, Any]:
|
|
with torch.no_grad():
|
|
if 'img' in input and input['img'] is not None:
|
|
input_img = input['img']
|
|
input_img = LoadImage.convert_to_img(input_img)
|
|
img_tensor = self.img_preprocessor(input_img)[None, ...]
|
|
|
|
if self.device_id >= 0:
|
|
img_tensor = img_tensor.to('cuda:{}'.format(
|
|
self.device_id))
|
|
_, _, image_feature, image_tensors = self.model.get_feature(
|
|
None, None, img_tensor)
|
|
image_feature = image_feature.cpu().numpy()
|
|
else:
|
|
image_feature, image_tensors = None, None
|
|
|
|
if 'text' in input and input['text'] is not None:
|
|
text_str = input['text']
|
|
if isinstance(text_str, str):
|
|
text_ids_tensor, text_mask_tensor = self.tokenize_text(
|
|
text_str)
|
|
else:
|
|
raise TypeError(
|
|
f'text should be str, but got {type(text_str)}')
|
|
|
|
if self.device_id >= 0:
|
|
text_ids_tensor = text_ids_tensor.to('cuda:{}'.format(
|
|
self.device_id))
|
|
text_mask_tensor = text_mask_tensor.to('cuda:{}'.format(
|
|
self.device_id))
|
|
text_feature, text_tensors, _, _ = self.model.get_feature(
|
|
text_ids_tensor, text_mask_tensor, None)
|
|
text_feature = text_feature.cpu().numpy()
|
|
else:
|
|
text_tensors, text_mask_tensor = None, None
|
|
|
|
if text_tensors is not None and text_mask_tensor is not None and image_tensors is not None:
|
|
score = self.model.get_cross_score(text_tensors,
|
|
text_mask_tensor,
|
|
image_tensors)[0].item()
|
|
else:
|
|
score = None
|
|
output = {
|
|
OutputKeys.IMG_EMBEDDING: image_feature,
|
|
OutputKeys.TEXT_EMBEDDING: text_feature,
|
|
OutputKeys.SCORES: score
|
|
}
|
|
return output
|
|
|
|
def postprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
|
|
return inputs
|