File size: 3,258 Bytes
04fbff5 |
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 73 74 75 76 77 |
import os
import json
import torch
import numpy as np
from tqdm import tqdm
from vbench.utils import load_video, load_dimension_info, read_frames_decord_by_fps
from vbench.third_party.grit_model import DenseCaptioning
import logging
logging.basicConfig(level = logging.INFO,format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
def get_dect_from_grit(model, image_arrays):
pred = []
if type(image_arrays) is not list and type(image_arrays) is not np.ndarray:
image_arrays = image_arrays.numpy()
with torch.no_grad():
for frame in image_arrays:
ret = model.run_caption_tensor(frame)
cur_pred = []
if len(ret[0])<1:
cur_pred.append(['',''])
else:
for idx, cap_det in enumerate(ret[0]):
cur_pred.append([cap_det[0], cap_det[2][0]])
pred.append(cur_pred)
return pred
def check_generate(color_key, object_key, predictions):
cur_object_color, cur_object = 0, 0
for frame_pred in predictions:
object_flag, color_flag = False, False
for pred in frame_pred:
if object_key == pred[1]:
for color_query in ["white","red","pink","blue","silver","purple","orange","green","gray","yellow","black","grey"]:
if color_query in pred[0]:
object_flag =True
if color_key in pred[0]:
color_flag = True
if color_flag:
cur_object_color+=1
if object_flag:
cur_object +=1
return cur_object, cur_object_color
def color(model, video_dict, device):
success_frame_count_all, video_count = 0, 0
video_results = []
for info in tqdm(video_dict):
if 'auxiliary_info' not in info:
raise "Auxiliary info is not in json, please check your json."
# print(info)
color_info = info['auxiliary_info']['color']
object_info = info['prompt']
object_info = object_info.replace('a ','').replace('an ','').replace(color_info,'').strip()
for video_path in info['video_list']:
video_arrays = load_video(video_path, num_frames=16, return_tensor=False)
cur_video_pred = get_dect_from_grit(model ,video_arrays)
cur_object, cur_object_color = check_generate(color_info, object_info, cur_video_pred)
if cur_object>0:
cur_success_frame_rate = cur_object_color/cur_object
success_frame_count_all += cur_success_frame_rate
video_count += 1
video_results.append({'video_path': video_path, 'video_results': cur_success_frame_rate})
success_rate = success_frame_count_all / video_count
return success_rate, video_results
def compute_color(json_dir, device, submodules_dict):
dense_caption_model = DenseCaptioning(device)
dense_caption_model.initialize_model(**submodules_dict)
logger.info("Initialize detection model success")
_, prompt_dict_ls = load_dimension_info(json_dir, dimension='color', lang='en')
all_results, video_results = color(dense_caption_model, prompt_dict_ls, device)
return all_results, video_results
|