def main():
if len(sys.argv) < 5:
print_usage()
return 2
try:
train_dataset_filepath = sys.argv[1]
feature_param_filepath = sys.argv[2]
model_param_filepath = sys.argv[3]
output_model_filepath = sys.argv[4]
except ValueError:
print_usage()
return 2
try:
train_dataset = import_python_file(train_dataset_filepath)
feature_param = import_python_file(feature_param_filepath)
model_param = import_python_file(model_param_filepath)
except Exception as e:
print("Error: %s" % e)
return 1
cache_result = cmd_option_exists(sys.argv, 3, len(sys.argv),
'--cache-result')
parallelize = cmd_option_exists(sys.argv, 3, len(sys.argv),
'--parallelize')
suppress_plot = cmd_option_exists(sys.argv, 3, len(sys.argv),
'--suppress-plot')
pool_method = get_cmd_option(sys.argv, 3, len(sys.argv), '--pool')
if not (pool_method is None or pool_method in POOL_METHODS):
print('--pool can only have option among {}'.format(
', '.join(POOL_METHODS)))
return 2
subj_model = get_cmd_option(sys.argv, 3, len(sys.argv), '--subj-model')
try:
if subj_model is not None:
from sureal.subjective_model import SubjectiveModel
subj_model_class = SubjectiveModel.find_subclass(subj_model)
else:
subj_model_class = None
except Exception as e:
print("Error: %s" % e)
return 1
save_plot_dir = get_cmd_option(sys.argv, 3, len(sys.argv), '--save-plot')
if cache_result:
result_store = FileSystemResultStore()
else:
result_store = None
# pooling
if pool_method == 'harmonic_mean':
aggregate_method = ListStats.harmonic_mean
elif pool_method == 'min':
aggregate_method = np.min
elif pool_method == 'median':
aggregate_method = np.median
elif pool_method == 'perc5':
aggregate_method = ListStats.perc5
elif pool_method == 'perc10':
aggregate_method = ListStats.perc10
elif pool_method == 'perc20':
aggregate_method = ListStats.perc20
else: # None or 'mean'
aggregate_method = np.mean
logger = None
try:
if suppress_plot:
raise AssertionError
from vmaf import plt
fig, ax = plt.subplots(figsize=(5, 5), nrows=1, ncols=1)
train_test_vmaf_on_dataset(
train_dataset=train_dataset,
test_dataset=None,
feature_param=feature_param,
model_param=model_param,
train_ax=ax,
test_ax=None,
result_store=result_store,
parallelize=parallelize,
logger=logger,
output_model_filepath=output_model_filepath,
aggregate_method=aggregate_method,
subj_model_class=subj_model_class,
)
bbox = {'facecolor': 'white', 'alpha': 0.5, 'pad': 20}
ax.annotate('Training Set',
xy=(0.1, 0.85),
xycoords='axes fraction',
bbox=bbox)
# ax.set_xlim([-10, 110])
# ax.set_ylim([-10, 110])
plt.tight_layout()
if save_plot_dir is None:
DisplayConfig.show()
else:
DisplayConfig.show(write_to_dir=save_plot_dir)
except ImportError:
print_matplotlib_warning()
train_test_vmaf_on_dataset(
train_dataset=train_dataset,
test_dataset=None,
feature_param=feature_param,
model_param=model_param,
train_ax=None,
test_ax=None,
result_store=result_store,
parallelize=parallelize,
logger=logger,
output_model_filepath=output_model_filepath,
aggregate_method=aggregate_method,
subj_model_class=subj_model_class,
)
except AssertionError:
train_test_vmaf_on_dataset(
train_dataset=train_dataset,
test_dataset=None,
feature_param=feature_param,
model_param=model_param,
train_ax=None,
test_ax=None,
result_store=result_store,
parallelize=parallelize,
logger=logger,
output_model_filepath=output_model_filepath,
aggregate_method=aggregate_method,
subj_model_class=subj_model_class,
)
return 0
def test_train_test_on_raw_dataset_with_dis1st_thr(self):
from vmaf.routine import train_test_vmaf_on_dataset
train_dataset = import_python_file(
VmafConfig.test_resource_path('raw_dataset_sample.py'))
model_param = import_python_file(
VmafConfig.test_resource_path('model_param_sample.py'))
feature_param = import_python_file(
VmafConfig.test_resource_path('feature_param_sample.py'))
train_fassembler, train_assets, train_stats, test_fassembler, test_assets, test_stats, _ = train_test_vmaf_on_dataset(
train_dataset=train_dataset,
test_dataset=train_dataset,
feature_param=feature_param,
model_param=model_param,
train_ax=None,
test_ax=None,
result_store=None,
parallelize=True,
logger=None,
fifo_mode=True,
output_model_filepath=self.output_model_filepath)
self.train_fassembler = train_fassembler
self.assertTrue(os.path.exists(self.output_model_filepath))
self.assertAlmostEqual(train_stats['ys_label_pred'][0],
93.565459224020742,
places=3)
self.assertAlmostEqual(test_stats['ys_label_pred'][0],
93.565459224020742,
places=3)
def test_train_test_on_dataset_with_dis1st_thr_with_feature_optional_dict_good(self):
from vmaf.routine import train_test_vmaf_on_dataset
train_dataset = import_python_file(
VmafConfig.test_resource_path('dataset_sample.py'))
model_param = import_python_file(
VmafConfig.test_resource_path('model_param_sample.py'))
feature_param = import_python_file(
VmafConfig.test_resource_path('feature_param_sample_with_optional_dict_good.py'))
train_fassembler, train_assets, train_stats, test_fassembler, test_assets, test_stats, _ = train_test_vmaf_on_dataset(
train_dataset=train_dataset,
test_dataset=train_dataset,
feature_param=feature_param,
model_param=model_param,
train_ax=None,
test_ax=None,
result_store=None,
parallelize=True,
logger=None,
fifo_mode=True,
output_model_filepath=self.output_model_filepath,
)
self.train_fassembler = train_fassembler
self.assertTrue(os.path.exists(self.output_model_filepath))
self.assertAlmostEqual(train_stats['ys_label_pred'][0], 90.753010402770798, places=3)
self.assertAlmostEqual(test_stats['ys_label_pred'][0], 90.753010402770798, places=3)
runner = VmafQualityRunner(
train_assets,
None, fifo_mode=True,
delete_workdir=True,
result_store=None,
optional_dict={'model_filepath': self.output_model_filepath}
)
runner.run(parallelize=True)
results = runner.results
self.assertAlmostEqual(results[0]['VMAF_score'], 89.55494473011981, places=4)
self.assertAlmostEqual(results[1]['VMAF_score'], 61.01289549048653, places=4)
self.assertAlmostEqual(results[2]['VMAF_score'], 90.75301241304798, places=4)
self.assertAlmostEqual(results[3]['VMAF_score'], 89.27013895870179, places=4)
def test_train_test_on_dataset_with_dis1st_thr_with_feature_optional_dict(self):
from vmaf.routine import train_test_vmaf_on_dataset
train_dataset = import_python_file(
VmafConfig.test_resource_path('dataset_sample.py'))
model_param = import_python_file(
VmafConfig.test_resource_path('model_param_sample.py'))
feature_param = import_python_file(
VmafConfig.test_resource_path('feature_param_sample_with_optional_dict.py'))
with self.assertRaises(AssertionError):
# adm_ref_display_height 108000 exceeds the maximum allowed
train_fassembler, train_assets, train_stats, test_fassembler, test_assets, test_stats, _ = train_test_vmaf_on_dataset(
train_dataset=train_dataset,
test_dataset=train_dataset,
feature_param=feature_param,
model_param=model_param,
train_ax=None,
test_ax=None,
result_store=None,
parallelize=False,
logger=None,
fifo_mode=True,
output_model_filepath=self.output_model_filepath,
)
def main():
if len(sys.argv) < 5:
print_usage()
return 2
try:
train_dataset_filepath = sys.argv[1]
feature_param_filepath = sys.argv[2]
model_param_filepath = sys.argv[3]
output_model_filepath = sys.argv[4]
except ValueError:
print_usage()
return 2
try:
train_dataset = import_python_file(train_dataset_filepath)
feature_param = import_python_file(feature_param_filepath)
model_param = import_python_file(model_param_filepath)
except Exception as e:
print "Error: " + str(e)
return 1
cache_result = cmd_option_exists(sys.argv, 3, len(sys.argv), '--cache-result')
parallelize = cmd_option_exists(sys.argv, 3, len(sys.argv), '--parallelize')
suppress_plot = cmd_option_exists(sys.argv, 3, len(sys.argv), '--suppress-plot')
pool_method = get_cmd_option(sys.argv, 3, len(sys.argv), '--pool')
if not (pool_method is None
or pool_method in POOL_METHODS):
print '--pool can only have option among {}'.format(', '.join(POOL_METHODS))
return 2
subj_model = get_cmd_option(sys.argv, 3, len(sys.argv), '--subj-model')
try:
if subj_model is not None:
subj_model_class = SubjectiveModel.find_subclass(subj_model)
else:
subj_model_class = None
except Exception as e:
print "Error: " + str(e)
return 1
if cache_result:
result_store = FileSystemResultStore()
else:
result_store = None
# pooling
if pool_method == 'harmonic_mean':
aggregate_method = ListStats.harmonic_mean
elif pool_method == 'min':
aggregate_method = np.min
elif pool_method == 'median':
aggregate_method = np.median
elif pool_method == 'perc5':
aggregate_method = ListStats.perc5
elif pool_method == 'perc10':
aggregate_method = ListStats.perc10
elif pool_method == 'perc20':
aggregate_method = ListStats.perc20
else: # None or 'mean'
aggregate_method = np.mean
logger = None
try:
if suppress_plot:
raise AssertionError
import matplotlib.pyplot as plt
fig, ax = plt.subplots(figsize=(5, 5), nrows=1, ncols=1)
train_test_vmaf_on_dataset(train_dataset=train_dataset, test_dataset=None,
feature_param=feature_param, model_param=model_param,
train_ax=ax, test_ax=None,
result_store=result_store,
parallelize=parallelize,
logger=logger,
output_model_filepath=output_model_filepath,
aggregate_method=aggregate_method,
subj_model_class=subj_model_class,
)
bbox = {'facecolor':'white', 'alpha':0.5, 'pad':20}
ax.annotate('Training Set', xy=(0.1, 0.85), xycoords='axes fraction', bbox=bbox)
# ax.set_xlim([-10, 110])
# ax.set_ylim([-10, 110])
plt.tight_layout()
plt.show()
except ImportError:
print_matplotlib_warning()
train_test_vmaf_on_dataset(train_dataset=train_dataset, test_dataset=None,
feature_param=feature_param, model_param=model_param,
train_ax=None, test_ax=None,
result_store=result_store,
parallelize=parallelize,
logger=logger,
output_model_filepath=output_model_filepath,
aggregate_method=aggregate_method,
subj_model_class=subj_model_class,
)
except AssertionError:
train_test_vmaf_on_dataset(train_dataset=train_dataset, test_dataset=None,
feature_param=feature_param, model_param=model_param,
train_ax=None, test_ax=None,
result_store=result_store,
parallelize=parallelize,
logger=logger,
output_model_filepath=output_model_filepath,
aggregate_method=aggregate_method,
subj_model_class=subj_model_class,
)
return 0
def test_train_test_on_dataset_with_dis1st_thr(self):
train_dataset = import_python_file(
config.ROOT + '/python/test/resource/dataset_sample.py')
model_param = import_python_file(
config.ROOT + '/python/test/resource/model_param_sample.py')
feature_param = import_python_file(
config.ROOT + '/python/test/resource/feature_param_sample.py')
train_fassembler, train_assets, train_stats, test_fassembler, test_assets, test_stats, _ = train_test_vmaf_on_dataset(
train_dataset=train_dataset,
test_dataset=train_dataset,
feature_param=feature_param,
model_param=model_param,
train_ax=None,
test_ax=None,
result_store=None,
parallelize=True,
logger=None,
fifo_mode=True,
output_model_filepath=self.output_model_filepath,
)
self.train_fassembler = train_fassembler
self.assertTrue(os.path.exists(self.output_model_filepath))
self.assertAlmostEqual(train_stats['ys_label_pred'][0],
90.753010402770798,
places=3)
self.assertAlmostEqual(test_stats['ys_label_pred'][0],
90.753010402770798,
places=3)
def test_train_test_on_raw_dataset_with_dis1st_thr(self):
train_dataset = import_python_file(
VmafConfig.test_resource_path('raw_dataset_sample.py'))
model_param = import_python_file(
VmafConfig.test_resource_path('model_param_sample.py'))
feature_param = import_python_file(
VmafConfig.test_resource_path('feature_param_sample.py'))
train_fassembler, train_assets, train_stats, test_fassembler, test_assets, test_stats, _ = train_test_vmaf_on_dataset(
train_dataset=train_dataset,
test_dataset=train_dataset,
feature_param=feature_param,
model_param=model_param,
train_ax=None,
test_ax=None,
result_store=None,
parallelize=True,
logger=None,
fifo_mode=True,
output_model_filepath=self.output_model_filepath
)
self.train_fassembler = train_fassembler
self.assertTrue(os.path.exists(self.output_model_filepath))
self.assertAlmostEqual(train_stats['ys_label_pred'][0], 93.565459224020742, places=3)
self.assertAlmostEqual(test_stats['ys_label_pred'][0], 93.565459224020742, places=3)
