zip_feat_and_scores.py

import numpy as np
from matplotlib import pyplot as plt 
import pandas as pd
import os
import glob
from joblib import load,dump

import argparse

parser = argparse.ArgumentParser(description='Zip ChipQA features from a folder with scores and names from a csv file into a single dictionary and save it')
parser.add_argument('input_folder',help='Folder containing input features generated by chipqa.py (format: name_of_video.z)')
parser.add_argument('csv_file',help='CSV file containing video name under column "video" and score under column "MOS"')
parser.add_argument('results_file',help='File where features and scores are stored. Must end with .z extension')

args = parser.parse_args()
dataset='apv'
names = []
score_list =[]
if(dataset=='apv'):
    csv_file =args.csv_file 
    score_csv = pd.read_csv(csv_file)

folder =args.input_folder 

filenames = glob.glob(os.path.join(folder,'*.z'))
feats = []
scores = []
ns = []
for i,file in enumerate(sorted(filenames)):
    fname = os.path.basename(file)
    print(fname)
    bname = os.path.splitext(fname)[0]
    yuv_fname = os.path.splitext(fname)[0]+'.mp4'

    if(dataset=='apv'):
         name = os.path.splitext(fname)[0]+'.yuv'
         score = score_csv[score_csv['video'] ==name].MOS.iloc[0]
    X = load(file)
    print(X)
    x = X['features']
    print(x.shape)
    y = score


    x = np.asarray(x)

    feats.append(x)
    scores.append(y)
    ns.append(fname)
feats = np.asarray(feats)
print(feats.shape)
X = {'features':feats,'score':scores,'name':ns}
dump(X,args.results_file,compress=('lzma',3))
print('dump done')
print()