|
from tqdm import tqdm, trange |
|
import torch |
|
import torch.nn.functional as F |
|
import numpy as np |
|
|
|
from utils.run_utils import topk_3d, generate_min_max_length_mask, extract_topk_elements |
|
from modules.ndcg_iou import calculate_ndcg_iou |
|
|
|
def grab_corpus_feature(model, corpus_loader, device): |
|
model.eval() |
|
all_video_feat, all_video_mask = [], [] |
|
all_sub_feat, all_sub_mask = [], [] |
|
|
|
|
|
with torch.no_grad(): |
|
for batch_input in tqdm(corpus_loader, desc="Compute Corpus Feature: ", total=len(corpus_loader)): |
|
batch_input = {k: v.to(device) for k, v in batch_input.items()} |
|
_video_feat, _sub_feat = model.encode_context(batch_input["video_feat"], batch_input["video_mask"], |
|
batch_input["sub_feat"], batch_input["sub_mask"]) |
|
|
|
all_video_feat.append(_video_feat.detach().cpu()) |
|
all_video_mask.append(batch_input["video_mask"].detach().cpu()) |
|
all_sub_feat.append(_sub_feat.detach().cpu()) |
|
all_sub_mask.append(batch_input["sub_mask"].detach().cpu()) |
|
|
|
all_video_feat = torch.cat(all_video_feat, dim=0) |
|
all_video_mask = torch.cat(all_video_mask, dim=0) |
|
all_sub_feat = torch.cat(all_sub_feat, dim=0) |
|
all_sub_mask = torch.cat(all_sub_mask, dim=0) |
|
|
|
return { "all_video_feat": all_video_feat, |
|
"all_video_mask": all_video_mask, |
|
"all_sub_feat": all_sub_feat, |
|
"all_sub_mask": all_sub_mask} |
|
|
|
|
|
def eval_epoch(model, corpus_feature, eval_loader, eval_gt, opt, corpus_video_list): |
|
topn_video = 100 |
|
device = opt.device |
|
model.eval() |
|
all_query_id = [] |
|
all_video_feat = corpus_feature["all_video_feat"].to(device) |
|
all_video_mask = corpus_feature["all_video_mask"].to(device) |
|
all_sub_feat = corpus_feature["all_sub_feat"].to(device) |
|
all_sub_mask = corpus_feature["all_sub_mask"].to(device) |
|
all_query_score, all_end_prob, all_start_prob, all_top_video_name = [], [], [], [] |
|
for batch_input in tqdm(eval_loader, desc="Compute Query Scores: ", total=len(eval_loader)): |
|
batch_input = {k: v.to(device) for k, v in batch_input.items()} |
|
query_scores, start_probs, end_probs = model.get_pred_from_raw_query( |
|
query_feat = batch_input["query_feat"], |
|
query_mask = batch_input["query_mask"], |
|
video_feat = all_video_feat, |
|
video_mask = all_video_mask, |
|
sub_feat = all_sub_feat, |
|
sub_mask = all_sub_mask, |
|
cross=True) |
|
query_scores = torch.exp(opt.q2c_alpha * query_scores) |
|
start_probs = F.softmax(start_probs, dim=-1) |
|
end_probs = F.softmax(end_probs, dim=-1) |
|
|
|
query_scores, start_probs, end_probs, video_name_top = extract_topk_elements(query_scores, start_probs, end_probs, corpus_video_list, topn_video) |
|
|
|
all_query_id.append(batch_input["query_id"].detach().cpu()) |
|
all_query_score.append(query_scores.detach().cpu()) |
|
all_start_prob.append(start_probs.detach().cpu()) |
|
all_end_prob.append(end_probs.detach().cpu()) |
|
all_top_video_name.extend(video_name_top) |
|
|
|
all_query_id = torch.cat(all_query_id, dim=0) |
|
all_query_id = all_query_id.tolist() |
|
|
|
all_query_score = torch.cat(all_query_score, dim=0) |
|
all_start_prob = torch.cat(all_start_prob, dim=0) |
|
all_end_prob = torch.cat(all_end_prob, dim=0) |
|
average_ndcg = calculate_average_ndcg(all_query_id, all_start_prob, all_query_score, all_end_prob, all_top_video_name, eval_gt, opt) |
|
return average_ndcg |
|
|
|
def calculate_average_ndcg(all_query_id, all_start_prob, all_query_score, all_end_prob, all_top_video_name, eval_gt, opt): |
|
topn_moment = max(opt.ndcg_topk) |
|
|
|
all_2D_map = torch.einsum("qvm,qv,qvn->qvmn", all_start_prob, all_query_score, all_end_prob) |
|
map_mask = generate_min_max_length_mask(all_2D_map.shape, min_l=opt.min_pred_l, max_l=opt.max_pred_l) |
|
all_2D_map = all_2D_map * map_mask |
|
all_pred = {} |
|
for idx in trange(len(all_2D_map), desc="Collect Predictions: "): |
|
query_id = all_query_id[idx] |
|
score_map = all_2D_map[idx] |
|
top_score, top_idx = topk_3d(score_map, topn_moment) |
|
top_video_name = all_top_video_name[idx] |
|
pred_videos = [top_video_name[i[0]] for i in top_idx] |
|
pre_start_time = [i[1].item() * opt.clip_length for i in top_idx] |
|
pre_end_time = [i[2].item() * opt.clip_length for i in top_idx] |
|
|
|
pred_result = [] |
|
for video_name, s, e, score, in zip(pred_videos, pre_start_time, pre_end_time, top_score): |
|
pred_result.append({ |
|
"video_name": video_name, |
|
"timestamp": [s, e], |
|
"model_scores": score |
|
}) |
|
|
|
all_pred[query_id] = pred_result |
|
|
|
average_ndcg = calculate_ndcg_iou(eval_gt, all_pred, opt.iou_threshold, opt.ndcg_topk) |
|
return average_ndcg |
