Merge pull request #60 from ai-forever/complex_filter

feat: add complex filter

Author: boomb0om

DPF/filters/__init__.py

@@ -1,2 +1,3 @@
 from .column_filter import ColumnFilter
+from .complex_filter import ComplexDataFilter
 from .data_filter import DataFilter

DPF/filters/complex_filter.py

@@ -0,0 +1,86 @@
+from dataclasses import dataclass
+from typing import Any
+
+from DPF.filters.data_filter import DataFilter
+from DPF.modalities import ModalityName
+from DPF.types import ModalityToDataMapping
+
+
+@dataclass
+class ComplexFilterPreprocessedData:
+    key: str
+    preprocessed_values: dict[int, Any]
+
+
+class ComplexDataFilter(DataFilter):
+
+    def __init__(
+        self,
+        datafilters: list[DataFilter],
+        workers: int,
+        pbar: bool = True,
+        _pbar_position: int = 0
+    ):
+        super().__init__(pbar, _pbar_position)
+        self.datafilters = datafilters
+        self.workers = workers
+
+        assert len(self.datafilters) > 0
+        assert all(
+            i.key_column == self.datafilters[0].key_column for i in self.datafilters
+        )  # check all filters have same key col
+
+    @property
+    def modalities(self) -> list[ModalityName]:
+        modals = []
+        for datafilter in self.datafilters:
+            modals.extend(datafilter.modalities)
+        return list(set(modals))
+
+    @property
+    def key_column(self) -> str:
+        return self.datafilters[0].key_column
+
+    @property
+    def metadata_columns(self) -> list[str]:
+        meta_cols = []
+        for datafilter in self.datafilters:
+            meta_cols.extend(datafilter.metadata_columns)
+        return list(set(meta_cols))
+
+    @property
+    def result_columns(self) -> list[str]:
+        result_cols = []
+        for datafilter in self.datafilters:
+            result_cols.extend(datafilter.result_columns)
+        return list(set(result_cols))
+
+    @property
+    def dataloader_kwargs(self) -> dict[str, Any]:
+        return {
+            "num_workers": self.workers,
+            "batch_size": 1,
+            "drop_last": False,
+        }
+
+    def preprocess_data(
+        self,
+        modality2data: ModalityToDataMapping,
+        metadata: dict[str, Any]
+    ) -> Any:
+        key = metadata[self.key_column]
+        preprocessed_results = {}
+        for ind, datafilter in enumerate(self.datafilters):
+            preprocessed_results[ind] = datafilter.preprocess_data(modality2data, metadata)
+        return ComplexFilterPreprocessedData(key, preprocessed_results)
+
+    def process_batch(self, batch: list[Any]) -> dict[str, list[Any]]:
+        results = {}
+        preprocessed_data: ComplexFilterPreprocessedData = batch[0]
+        for ind, datafilter in enumerate(self.datafilters):
+            filter_results_data = datafilter.process_batch([preprocessed_data.preprocessed_values[ind]])
+            for col, value in filter_results_data.items():
+                results[col] = value
+        results[self.key_column] = [preprocessed_data.key]
+
+        return results

pyproject.toml

@@ -18,7 +18,7 @@ dependencies = [
   "numpy",
   "soundfile",
   "scipy",
-  "pillow",
+  "pillow==10.3.0",
   "tqdm",
   "pandas",
   "pandarallel",

tests/test_complex_filter.py

@@ -0,0 +1,29 @@
+from DPF import DatasetReader
+from DPF.configs import ShardsDatasetConfig
+from DPF.filters import ComplexDataFilter
+from DPF.filters.images.hash_filters import PHashFilter
+from DPF.filters.images.info_filter import ImageInfoFilter
+
+
+def test_shards_complex_filter():
+    path = 'tests/datasets/shards_correct'
+    config = ShardsDatasetConfig.from_path_and_columns(
+        path,
+        image_name_col="image_name",
+        text_col="caption"
+    )
+
+    reader = DatasetReader()
+    dataset = reader.read_from_config(config)
+    phashfilter = PHashFilter(workers=1)
+    infofilter = ImageInfoFilter(workers=1)
+
+    datafilter = ComplexDataFilter([phashfilter, infofilter], workers=2)
+    dataset.apply_data_filter(datafilter)
+
+    assert not dataset.df['image_phash_8'].isna().any()
+    assert dataset.df['is_correct'].all()
+    assert dataset.df['error'].isna().all()
+    assert not dataset.df['width'].isna().any()
+    assert not dataset.df['height'].isna().any()
+    assert not dataset.df['channels'].isna().any()