| 22 | |
| 23 | |
| 24 | class StrredFeatureExtractor(MatlabFeatureExtractor): |
| 25 | TYPE = 'STRRED_feature' |
| 26 | |
| 27 | # VERSION = '1.0' |
| 28 | # VERSION = '1.1' # fix matlab code where width and height are mistakenly swapped |
| 29 | # VERSION = '1.2' # fix minor frame and prev frame swap issue |
| 30 | VERSION = '1.3' # align ST-RRED with ST-RREDopt calculations |
| 31 | |
| 32 | ATOM_FEATURES = ['srred', 'trred', ] |
| 33 | |
| 34 | DERIVED_ATOM_FEATURES = ['strred', ] |
| 35 | |
| 36 | MATLAB_WORKSPACE = VmafConfig.root_path('matlab', 'strred') |
| 37 | |
| 38 | @classmethod |
| 39 | def _assert_an_asset(cls, asset): |
| 40 | super(StrredFeatureExtractor, cls)._assert_an_asset(asset) |
| 41 | assert asset.ref_yuv_type == 'yuv420p' and asset.dis_yuv_type == 'yuv420p', \ |
| 42 | 'STRRED feature extractor only supports yuv420p for now.' |
| 43 | |
| 44 | def _generate_result(self, asset): |
| 45 | # routine to call the command-line executable and generate quality |
| 46 | # scores in the log file. |
| 47 | |
| 48 | ref_procfile_path = asset.ref_procfile_path |
| 49 | dis_procfile_path = asset.dis_procfile_path |
| 50 | log_file_path = self._get_log_file_path(asset) |
| 51 | |
| 52 | current_dir = os.getcwd() + '/' |
| 53 | |
| 54 | ref_procfile_path = make_absolute_path(ref_procfile_path, current_dir) |
| 55 | dis_procfile_path = make_absolute_path(dis_procfile_path, current_dir) |
| 56 | log_file_path = make_absolute_path(log_file_path, current_dir) |
| 57 | |
| 58 | quality_width, quality_height = asset.quality_width_height |
| 59 | |
| 60 | strred_cmd = '''{matlab} -nodisplay -nosplash -nodesktop -r "run_strred('{ref}', '{dis}', {h}, {w}); exit;" >> {log_file_path}'''.format( |
| 61 | matlab=VmafExternalConfig.get_and_assert_matlab(), |
| 62 | ref=ref_procfile_path, |
| 63 | dis=dis_procfile_path, |
| 64 | w=quality_width, |
| 65 | h=quality_height, |
| 66 | log_file_path=log_file_path, |
| 67 | ) |
| 68 | if self.logger: |
| 69 | self.logger.info(strred_cmd) |
| 70 | |
| 71 | os.chdir(self.MATLAB_WORKSPACE) |
| 72 | run_process(strred_cmd, shell=True) |
| 73 | os.chdir(current_dir) |
| 74 | |
| 75 | @classmethod |
| 76 | @override(Executor) |
| 77 | def _post_process_result(cls, result): |
| 78 | |
| 79 | def _strred(srred_trred): |
| 80 | srred, trred = srred_trred |
| 81 | try: |