File size: 3,264 Bytes
211b431 |
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 pandas as pd
from scipy.io import loadmat, savemat
import numpy as np
def split_features(data_name, df, result_file, feature_path, layer_name):
data = loadmat(result_file)
all_variables = {}
for key, value in data.items():
# ignore '__'
if not key.startswith('__') and not key.endswith('__'):
all_variables[key] = value
if data_name == 'konvid_1k':
for i in range(len(all_variables['Test_videos_Median_model'])):
test_vids = all_variables['Test_videos_Median_model'][i]
test_vids = test_vids.tolist()
else:
test_vids = []
for i in range(len(all_variables['Test_videos_Median_model'])):
vid = all_variables['Test_videos_Median_model'][i].strip()
test_vids.append(vid)
# drop greyscale videos
if data_name == 'youtube_ugc':
grey_df = pd.read_csv(f'{metadata_path}/greyscale_report/{data_name.upper()}_greyscale_metadata.csv')
grey_indices = grey_df.iloc[:, 0].tolist()
df = df.drop(index=grey_indices).reset_index(drop=True)
all_vids = df.iloc[:, 0].tolist()
print(all_vids)
print(test_vids)
train_vids = list(set(all_vids) - set(test_vids))
print(len(test_vids))
print(len(train_vids))
# split all_data into train and test based on vids
train_df = df[df.iloc[:, 0].isin(train_vids)]
test_df = df[df.iloc[:, 0].isin(test_vids)]
print(len(test_df))
# reorder columns
sorted_train_df = pd.DataFrame({'vid': train_df.iloc[:, 0], 'framerate': train_df['framerate'], 'MOS': train_df['mos']})
sorted_test_df = pd.DataFrame({'vid': test_df.iloc[:, 0], 'framerate': test_df['framerate'], 'MOS': test_df['mos']})
# use indices from the train and test DataFrames to split features
data = loadmat(f'{feature_path}{layer_name}/original_features/{network_name}_{data_name}_original_features.mat')
features = data[f'{data_name}']
if data_name == 'youtube_ugc':
features = np.delete(features, grey_indices, axis=0)
train_features = features[train_df.index]
test_features = features[test_df.index]
# recover the median train and test files
sorted_train_df.to_csv(f'{metadata_path}mos_files/{data_name}_MOS_train.csv', index=False)
sorted_test_df.to_csv(f'{metadata_path}mos_files/{data_name}_MOS_test.csv', index=False)
savemat(f'{feature_path}{layer_name}/relaxvqa_{data_name}_original_train_features.mat', {f'{data_name}_train_features': train_features})
savemat(f'{feature_path}{layer_name}/relaxvqa_{data_name}_original_test_features.mat', {f'{data_name}_test_features': test_features})
return train_features, test_features, test_vids
if __name__ == '__main__':
metadata_path = '../../metadata/'
feature_path = '../../features_merged_frag/'
result_path = f'../../log/result/'
data_name = 'cvd_2014'
network_name = 'relaxvqa'
layer_name = 'pool'
model_name = 'Mlp'
select_criteria = 'byrmse'
df = pd.read_csv(f'{metadata_path}/{data_name.upper()}_metadata.csv')
result_file = f'{result_path}{data_name}_{network_name}_{select_criteria}.mat'
train_features, test_features, test_vids = split_features(data_name, df, result_file, feature_path, layer_name)
|