-
Notifications
You must be signed in to change notification settings - Fork 195
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
* add new OP: audio_size_filter * rename get_image_size to get_file_size
- Loading branch information
Showing
9 changed files
with
221 additions
and
8 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,74 @@ | ||
| import numpy as np | ||
|
|
||
| from data_juicer.utils.constant import Fields, StatsKeys | ||
| from data_juicer.utils.mm_utils import get_file_size, size_to_bytes | ||
|
|
||
| from ..base_op import OPERATORS, Filter | ||
|
|
||
|
|
||
| @OPERATORS.register_module('audio_size_filter') | ||
| class AudioSizeFilter(Filter): | ||
| """Keep data samples whose audio size (in bytes/kb/MB/...) within a | ||
| specific range. | ||
| """ | ||
|
|
||
| def __init__(self, | ||
| min_size: str = '0', | ||
| max_size: str = '1TB', | ||
| any_or_all: str = 'any', | ||
| *args, | ||
| **kwargs): | ||
| """ | ||
| Initialization method. | ||
| :param min_size: The min audio size to keep samples. set to be "0" by | ||
| default for no size constraint | ||
| :param max_size: The max audio size to keep samples. set to be | ||
| "1Tb" by default, an approximate for un-limited case | ||
| :param any_or_all: keep this sample with 'any' or 'all' strategy of | ||
| all audios. 'any': keep this sample if any audios meet the | ||
| condition. 'all': keep this sample only if all audios meet the | ||
| condition. | ||
| :param args: extra args | ||
| :param kwargs: extra args | ||
| """ | ||
| super().__init__(*args, **kwargs) | ||
| self.min_size = size_to_bytes(min_size) | ||
| self.max_size = size_to_bytes(max_size) | ||
| if any_or_all not in ['any', 'all']: | ||
| raise ValueError(f'Keep strategy [{any_or_all}] is not supported. ' | ||
| f'Can only be one of ["any", "all"].') | ||
| self.any = (any_or_all == 'any') | ||
|
|
||
| def compute_stats(self, sample, context=False): | ||
| # check if it's computed already | ||
| if StatsKeys.audio_sizes in sample[Fields.stats]: | ||
| return sample | ||
|
|
||
| # there is no audio in this sample | ||
| if self.audio_key not in sample or not sample[self.audio_key]: | ||
| sample[Fields.stats][StatsKeys.audio_sizes] = np.array( | ||
| [], dtype=np.float64) | ||
| return sample | ||
|
|
||
| # for size calculation, no need to load audios into memory | ||
| sample[Fields.stats][StatsKeys.audio_sizes] = [ | ||
| get_file_size(aud_path) for aud_path in sample[self.audio_key] | ||
| ] | ||
|
|
||
| return sample | ||
|
|
||
| def process(self, sample): | ||
| audio_sizes = sample[Fields.stats][StatsKeys.audio_sizes] | ||
| keep_bools = np.array([ | ||
| self.min_size <= audio_size <= self.max_size | ||
| for audio_size in audio_sizes | ||
| ]) | ||
| if len(keep_bools) <= 0: | ||
| return True | ||
|
|
||
| # different strategies | ||
| if self.any: | ||
| return keep_bools.any() | ||
| else: | ||
| return keep_bools.all() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,132 @@ | ||
| import os | ||
| import unittest | ||
|
|
||
| from datasets import Dataset | ||
|
|
||
| from data_juicer.ops.filter.audio_size_filter import AudioSizeFilter | ||
| from data_juicer.utils.constant import Fields | ||
|
|
||
|
|
||
| class AudioSizeFilterTest(unittest.TestCase): | ||
|
|
||
| data_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), | ||
| '..', 'data') | ||
| aud1_path = os.path.join(data_path, 'audio1.wav') # 970574 / 948K | ||
| aud2_path = os.path.join(data_path, 'audio2.wav') # 2494872 / 2.4M | ||
| aud3_path = os.path.join(data_path, 'audio3.ogg') # 597254 / 583K | ||
|
|
||
| def _run_audio_size_filter(self,dataset: Dataset, target_list, op, np=1): | ||
| if Fields.stats not in dataset.features: | ||
| dataset = dataset.add_column(name=Fields.stats, | ||
| column=[{}] * dataset.num_rows) | ||
| dataset = dataset.map(op.compute_stats, num_proc=np) | ||
| dataset = dataset.filter(op.process, num_proc=np) | ||
| dataset = dataset.select_columns(column_names=[op.audio_key]) | ||
| res_list = dataset.to_list() | ||
| self.assertEqual(res_list, target_list) | ||
|
|
||
| def test_min_max(self): | ||
|
|
||
| ds_list = [{ | ||
| 'audios': [self.aud1_path] | ||
| }, { | ||
| 'audios': [self.aud2_path] | ||
| }, { | ||
| 'audios': [self.aud3_path] | ||
| }] | ||
| tgt_list = [{ | ||
| 'audios': [self.aud1_path] | ||
| }] | ||
| dataset = Dataset.from_list(ds_list) | ||
| op = AudioSizeFilter(min_size="800kb", max_size="1MB") | ||
| self._run_audio_size_filter(dataset, tgt_list, op) | ||
|
|
||
| def test_min(self): | ||
|
|
||
| ds_list = [{ | ||
| 'audios': [self.aud1_path] | ||
| }, { | ||
| 'audios': [self.aud2_path] | ||
| }, { | ||
| 'audios': [self.aud3_path] | ||
| }] | ||
| tgt_list = [{ | ||
| 'audios': [self.aud1_path] | ||
| }, { | ||
| 'audios': [self.aud2_path] | ||
| }] | ||
| dataset = Dataset.from_list(ds_list) | ||
| op = AudioSizeFilter(min_size="900kib") | ||
| self._run_audio_size_filter(dataset, tgt_list, op) | ||
|
|
||
| def test_max(self): | ||
|
|
||
| ds_list = [{ | ||
| 'audios': [self.aud1_path] | ||
| }, { | ||
| 'audios': [self.aud2_path] | ||
| }, { | ||
| 'audios': [self.aud3_path] | ||
| }] | ||
| tgt_list = [{ | ||
| 'audios': [self.aud1_path] | ||
| }, { | ||
| 'audios': [self.aud3_path] | ||
| }] | ||
| dataset = Dataset.from_list(ds_list) | ||
| op = AudioSizeFilter(max_size="2MiB") | ||
| self._run_audio_size_filter(dataset, tgt_list, op) | ||
|
|
||
| def test_any(self): | ||
|
|
||
| ds_list = [{ | ||
| 'audios': [self.aud1_path, self.aud2_path] | ||
| }, { | ||
| 'audios': [self.aud2_path, self.aud3_path] | ||
| }, { | ||
| 'audios': [self.aud1_path, self.aud3_path] | ||
| }] | ||
| tgt_list = [{ | ||
| 'audios': [self.aud1_path, self.aud2_path] | ||
| }, { | ||
| 'audios': [self.aud1_path, self.aud3_path] | ||
| }] | ||
| dataset = Dataset.from_list(ds_list) | ||
| op = AudioSizeFilter(min_size="800kb", max_size="1MB", | ||
| any_or_all='any') | ||
| self._run_audio_size_filter(dataset, tgt_list, op) | ||
|
|
||
| def test_all(self): | ||
|
|
||
| ds_list = [{ | ||
| 'audios': [self.aud1_path, self.aud2_path] | ||
| }, { | ||
| 'audios': [self.aud2_path, self.aud3_path] | ||
| }, { | ||
| 'audios': [self.aud1_path, self.aud3_path] | ||
| }] | ||
| tgt_list = [] | ||
| dataset = Dataset.from_list(ds_list) | ||
| op = AudioSizeFilter(min_size="800kb", max_size="1MB", | ||
| any_or_all='all') | ||
| self._run_audio_size_filter(dataset, tgt_list, op) | ||
|
|
||
| def test_filter_in_parallel(self): | ||
|
|
||
| ds_list = [{ | ||
| 'audios': [self.aud1_path] | ||
| }, { | ||
| 'audios': [self.aud2_path] | ||
| }, { | ||
| 'audios': [self.aud3_path] | ||
| }] | ||
| tgt_list = [{ | ||
| 'audios': [self.aud1_path] | ||
| }] | ||
| dataset = Dataset.from_list(ds_list) | ||
| op = AudioSizeFilter(min_size="800kb", max_size="1MB") | ||
| self._run_audio_size_filter(dataset, tgt_list, op, np=2) | ||
|
|
||
|
|
||
| if __name__ == '__main__': | ||
| unittest.main() |