| """ |
| Copyright (c) 2024, Alliance for Open Media. All rights reserved |
| |
| This source code is subject to the terms of the BSD 3-Clause Clear License |
| and the Alliance for Open Media Patent License 1.0. If the BSD 3-Clause Clear |
| License was not distributed with this source code in the LICENSE file, you |
| can obtain it at aomedia.org/license/software-license/bsd-3-c-c/. If the |
| Alliance for Open Media Patent License 1.0 was not distributed with this |
| source code in the PATENTS file, you can obtain it at |
| aomedia.org/license/patent-license/. |
| """ |
| import os |
| |
| |
| class FileCollector: |
| def __init__( |
| self, |
| data_path, |
| file_extension, |
| ctx_group="", |
| coding_config="", |
| desired_qp=(), |
| subtext="", |
| starttext="", |
| ): |
| # non-public |
| self._data_path = data_path |
| self._file_extension = "." + file_extension # e.g., 'csv' and 'json' |
| self._ctx_group = ctx_group |
| self._coding_config = coding_config |
| self._desired_qp = desired_qp |
| self._subtext = subtext |
| self._starttext = starttext |
| # public |
| self.all_files = self._get_all_files() |
| self.files = self._filter_files() |
| |
| # public |
| def get_files(self): |
| return self.files |
| |
| # non-public |
| def _get_all_files(self): |
| files_all = os.listdir(self._data_path) |
| files = [f for f in files_all if f.endswith(self._file_extension)] |
| return sorted(files) |
| |
| def _filter_files(self): |
| filtered_files = self.all_files |
| # ctx_group |
| ctx_group = self._ctx_group |
| if ctx_group != "": |
| filtered_files = [ |
| f for f in filtered_files if "Stat_" + ctx_group + "_Bin_" in f |
| ] |
| # coding config (AI, RA, LD, etc..) |
| cfg = self._coding_config |
| if cfg != "": |
| filtered_files = [f for f in filtered_files if "_" + cfg + "_" in f] |
| # filter based on desired QP list |
| desired_qp = self._desired_qp |
| if len(desired_qp) > 0: |
| all_filt_files = filtered_files.copy() |
| filtered_files = [] |
| for qp in desired_qp: |
| new_files = [f for f in all_filt_files if "_QP_" + str(qp) + "_" in f] |
| for f in new_files: |
| filtered_files.append(f) |
| # any subtext between underscores in '_{subtext}_' format |
| subtext = self._subtext |
| if subtext != "": |
| filtered_files = [f for f in filtered_files if subtext in f] |
| |
| starttext = self._starttext |
| if starttext != "": |
| filtered_files = [f for f in filtered_files if f.startswith(starttext)] |
| |
| return filtered_files |
