Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,614 Bytes
2d9a728 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 |
from .simple_tokenizer import SimpleTokenizer as _Tokenizer
from .viclip import ViCLIP
import torch
import numpy as np
import cv2
import os
def get_viclip(size='l',
pretrain=os.path.join(os.path.dirname(os.path.abspath(__file__)), "ViClip-InternVid-10M-FLT.pth")):
tokenizer = _Tokenizer()
vclip = ViCLIP(tokenizer=tokenizer, size=size, pretrain=pretrain)
m = {'viclip':vclip, 'tokenizer':tokenizer}
return m
def get_text_feat_dict(texts, clip, tokenizer, text_feat_d={}):
for t in texts:
feat = clip.get_text_features(t, tokenizer, text_feat_d)
text_feat_d[t] = feat
return text_feat_d
def get_vid_feat(frames, clip):
return clip.get_vid_features(frames)
def _frame_from_video(video):
while video.isOpened():
success, frame = video.read()
if success:
yield frame
else:
break
v_mean = np.array([0.485, 0.456, 0.406]).reshape(1,1,3)
v_std = np.array([0.229, 0.224, 0.225]).reshape(1,1,3)
def normalize(data):
return (data/255.0-v_mean)/v_std
def frames2tensor(vid_list, fnum=8, target_size=(224, 224), device=torch.device('cuda')):
assert(len(vid_list) >= fnum)
step = len(vid_list) // fnum
vid_list = vid_list[::step][:fnum]
vid_list = [cv2.resize(x[:,:,::-1], target_size) for x in vid_list]
vid_tube = [np.expand_dims(normalize(x), axis=(0, 1)) for x in vid_list]
vid_tube = np.concatenate(vid_tube, axis=1)
vid_tube = np.transpose(vid_tube, (0, 1, 4, 2, 3))
vid_tube = torch.from_numpy(vid_tube).to(device, non_blocking=True).float()
return vid_tube
def retrieve_text(frames,
texts,
models={'viclip':None,
'tokenizer':None},
topk=5,
device=torch.device('cuda')):
# clip, tokenizer = get_clip(name, model_cfg['size'], model_cfg['pretrained'], model_cfg['reload'])
assert(type(models)==dict and models['viclip'] is not None and models['tokenizer'] is not None)
clip, tokenizer = models['viclip'], models['tokenizer']
clip = clip.to(device)
frames_tensor = frames2tensor(frames, device=device)
vid_feat = get_vid_feat(frames_tensor, clip)
text_feat_d = {}
text_feat_d = get_text_feat_dict(texts, clip, tokenizer, text_feat_d)
text_feats = [text_feat_d[t] for t in texts]
text_feats_tensor = torch.cat(text_feats, 0)
probs, idxs = clip.get_predict_label(vid_feat, text_feats_tensor, top=topk)
ret_texts = [texts[i] for i in idxs.numpy()[0].tolist()]
return ret_texts, probs.numpy()[0] |