evaluate_one_dataset.py

import argparse 
import os 
import csv 
 
import torch 
 
import pandas as pd 
import numpy as np 
import pickle as pkl 
import decord 
import yaml 
 
from scipy import stats 
from sklearn.metrics import mean_squared_error 
from scipy.optimize import curve_fit 
from cover.datasets import UnifiedFrameSampler, spatial_temporal_view_decomposition 
from cover.models import COVER 
 
# use case 
# python evaluate_on_ytugc.py -o cover.yml -d cuda:3 --output result.csv -uh 0 
 
def save_to_csv(video_name, pre_smos, pre_tmos, pre_amos, pre_overall, filename): 
 combined_data = list(zip(video_name, pre_smos, pre_tmos, pre_amos, pre_overall)) 
 
 with open(filename, 'w', newline='') as csvfile: 
 writer = csv.writer(csvfile) 
 writer.writerow(['Video', 'semantic score', 'technical score', 'aesthetic score', 'overall/final score']) 
 writer.writerows(combined_data) 
 
mean_cover, std_cover = ( 
 torch.FloatTensor([123.675, 116.28, 103.53]), 
 torch.FloatTensor([58.395, 57.12, 57.375]), 
) 
 
mean_clip, std_clip = ( 
 torch.FloatTensor([122.77, 116.75, 104.09]), 
 torch.FloatTensor([68.50, 66.63, 70.32]) 
) 
 
def fuse_results(results: list): 
 x = (results[0] + results[1] + results[2]) 
 return { 
 "semantic" : results[0], 
 "technical": results[1], 
 "aesthetic": results[2], 
 "overall" : x, 
 } 
 
 
def gaussian_rescale(pr): 
 # The results should follow N(0,1) 
 pr = (pr - np.mean(pr)) / np.std(pr) 
 return pr 
 
 
def uniform_rescale(pr): 
 # The result scores should follow U(0,1) 
 return np.arange(len(pr))[np.argsort(pr).argsort()] / len(pr) 
 
 
def parse_args(): 
 parser = argparse.ArgumentParser() 
 parser.add_argument("-o", "--opt" , type=str, default="./cover.yml", help="the option file") 
 parser.add_argument('-d', "--device", type=str, default="cuda:0" , help='CUDA device id') 
 parser.add_argument("-t", "--target_set", type=str, default="val-ytugc", help="target_set") 
 parser.add_argument( "--output", type=str, default="ytugc.csv" , help='output file to store predict mos value') 
 args = parser.parse_args() 
 return args 
 
 
def logistic_func(X, bayta1, bayta2, bayta3, bayta4): 
 # 4-parameter logistic function 
 logisticPart = 1 + np.exp(np.negative(np.divide(X - bayta3, np.abs(bayta4)))) 
 yhat = bayta2 + np.divide(bayta1 - bayta2, logisticPart) 
 return yhat 
 
 
if __name__ == '__main__': 
 args = parse_args() 
 
 with open(args.opt, "r") as f: 
 opt = yaml.safe_load(f) 
 
 ### Load COVER 
 evaluator = COVER(**opt["model"]["args"]).to(args.device) 
 state_dict = torch.load(opt["test_load_path"], map_location=args.device) 
 
 # set strict=False here to avoid error of missing 
 # weight of prompt_learner in clip-iqa+, cross-gate 
 evaluator.load_state_dict(state_dict['state_dict'], strict=False) 
 
 dopt = opt["data"][args.target_set]["args"] 
 temporal_samplers = {} 
 for stype, sopt in dopt["sample_types"].items(): 
 temporal_samplers[stype] = UnifiedFrameSampler( 
 sopt["clip_len"] // sopt["t_frag"], 
 sopt["t_frag"], 
 sopt["frame_interval"], 
 sopt["num_clips"], 
 ) 
 
 if args.target_set == 'val-livevqc': 
 videos_dir = './datasets/LIVE_VQC/Video/' 
 datainfo = './datasets/LIVE_VQC/metainfo/LIVE_VQC_metadata.csv' 
 df = pd.read_csv(datainfo) 
 files = df['File'].tolist() 
 mos = df['MOS'].tolist() 
 elif args.target_set == 'val-kv1k': 
 videos_dir = './datasets/KoNViD/KoNViD_1k_videos/' 
 datainfo = './datasets/KoNViD/metainfo/KoNVid_metadata.csv' 
 df = pd.read_csv(datainfo) 
 files = df['Filename'].tolist() 
 files = [str(file) + '.mp4' for file in files] 
 mos = df['MOS'].tolist() 
 elif args.target_set == 'val-ytugc': 
 videos_dir = './datasets/YouTubeUGC/' 
 datainfo = './datasets/YouTubeUGC/../meta_info/Youtube-UGC_metadata.csv' 
 df = pd.read_csv(datainfo) 
 files = df['filename'].tolist() 
 mos = df['MOSFull'].tolist() 
 files = [str(file) + '_crf_10_ss_00_t_20.0.mp4' for file in files] 
 else: 
 print("unsupported video dataset for evaluation") 
 assert(0) 
 
 print(len(files)) 
 
 pure_name_list = [] 
 pre_overall = np.zeros(len(mos)) 
 pre_smos = np.zeros(len(mos)) 
 pre_tmos = np.zeros(len(mos)) 
 pre_amos = np.zeros(len(mos)) 
 gt_mos = np.array(mos) 
 count = 0 
 
 for vi in range(len(mos)): 
 video = files[vi] 
 pure_name = os.path.splitext(video)[0] 
 video_path = os.path.join(videos_dir, video) 
 
 views, _ = spatial_temporal_view_decomposition( 
 video_path, dopt["sample_types"], temporal_samplers 
 ) 
 
 for k, v in views.items(): 
 num_clips = dopt["sample_types"][k].get("num_clips", 1) 
 if k == 'technical' or k == 'aesthetic': 
 views[k] = ( 
 ((v.permute(1, 2, 3, 0) - mean_cover) / std_cover) 
 .permute(3, 0, 1, 2) 
 .reshape(v.shape[0], num_clips, -1, *v.shape[2:]) 
 .transpose(0, 1) 
 .to(args.device) 
 ) 
 elif k == 'semantic': 
 views[k] = ( 
 ((v.permute(1, 2, 3, 0) - mean_clip) / std_clip) 
 .permute(3, 0, 1, 2) 
 .reshape(v.shape[0], num_clips, -1, *v.shape[2:]) 
 .transpose(0, 1) 
 .to(args.device) 
 ) 
 
 results = [r.mean().item() for r in evaluator(views)] 
 
 
 pre_overall[count] = fuse_results(results)['overall'] 
 pre_smos[count] = results[0] 
 pre_tmos[count] = results[1] 
 pre_amos[count] = results[2] 
 pure_name_list.append(pure_name) 
 print("Process ", video, ", predicted quality score is ", pre_overall[count]) 
 count += 1 
 
 
 SROCC = stats.spearmanr(pre_overall, gt_mos)[0] 
 KROCC = stats.stats.kendalltau(pre_overall, gt_mos)[0] 
 
 # logistic regression btw y_pred & y 
 beta_init = [np.max(gt_mos), np.min(gt_mos), np.mean(pre_overall), 0.5] 
 popt, _ = curve_fit(logistic_func, pre_overall, gt_mos, p0=beta_init, maxfev=int(1e8)) 
 pre_overall_logistic = logistic_func(pre_overall, *popt) 
 
 PLCC = stats.pearsonr(gt_mos, pre_overall_logistic)[0] 
 RMSE = np.sqrt(mean_squared_error(gt_mos, pre_overall_logistic)) 
 
 print("Test results: SROCC={:.4f}, KROCC={:.4f}, PLCC={:.4f}, RMSE={:.4f}" 
 .format(SROCC, KROCC, PLCC, RMSE)) 
 
 save_to_csv(pure_name_list, pre_smos, pre_tmos, pre_amos, pre_overall, args.output)