|
import os |
|
import json |
|
|
|
import torch |
|
import numpy as np |
|
from tqdm import tqdm |
|
from vbench.utils import load_video, load_dimension_info |
|
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: |
|
image_arrays = image_arrays.numpy() |
|
with torch.no_grad(): |
|
for frame in image_arrays: |
|
try: |
|
pred.append(set(model.run_caption_tensor(frame)[0][0][2])) |
|
except: |
|
pred.append(set()) |
|
return pred |
|
|
|
def check_generate(key_info, predictions): |
|
cur_cnt = 0 |
|
for pred in predictions: |
|
if key_info in pred: |
|
cur_cnt+=1 |
|
return cur_cnt |
|
|
|
def object_class(model, video_dict, device): |
|
success_frame_count, frame_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." |
|
object_info = info['auxiliary_info']['object'] |
|
for video_path in info['video_list']: |
|
video_tensor = load_video(video_path, num_frames=16) |
|
cur_video_pred = get_dect_from_grit(model, video_tensor.permute(0,2,3,1)) |
|
cur_success_frame_count = check_generate(object_info, cur_video_pred) |
|
cur_success_frame_rate = cur_success_frame_count/len(cur_video_pred) |
|
success_frame_count += cur_success_frame_count |
|
frame_count += len(cur_video_pred) |
|
video_results.append({'video_path': video_path, 'video_results': cur_success_frame_rate}) |
|
success_rate = success_frame_count / frame_count |
|
return success_rate, video_results |
|
|
|
|
|
def compute_object_class(json_dir, device, submodules_dict): |
|
dense_caption_model = DenseCaptioning(device) |
|
dense_caption_model.initialize_model_det(**submodules_dict) |
|
logger.info("Initialize detection model success") |
|
_, prompt_dict_ls = load_dimension_info(json_dir, dimension='object_class', lang='en') |
|
all_results, video_results = object_class(dense_caption_model, prompt_dict_ls, device) |
|
return all_results, video_results |
|
|