feat: add filter pipelines

DPF/filters/images/dummy_gpu_filter.py

@@ -0,0 +1,53 @@
+from typing import Any, Union
+
+import torch
+
+from DPF.types import ModalityToDataMapping
+
+from .img_filter import ImageFilter
+
+
+class DummyGPUFilter(ImageFilter):
+ """
+ DummyGPUFilter class for testing purposes
+ """
+
+ def __init__(
+ self,
+ workers: int = 16,
+ device: Union[str, torch.device] = "cuda",
+ pbar: bool = True,
+ _pbar_position: int = 0
+ ):
+ super().__init__(pbar, _pbar_position)
+ self.num_workers = workers
+ self.device = device
+
+ @property
+ def result_columns(self) -> list[str]:
+ return ["dummy_label",]
+
+ @property
+ def dataloader_kwargs(self) -> dict[str, Any]:
+ return {
+ "num_workers": self.num_workers,
+ "batch_size": 1,
+ "drop_last": False,
+ }
+
+ def preprocess_data(
+ self,
+ modality2data: ModalityToDataMapping,
+ metadata: dict[str, Any]
+ ) -> Any:
+ key = metadata[self.key_column]
+ return key, 1
+
+ def process_batch(self, batch: list[Any]) -> dict[str, list[Any]]:
+ df_batch_labels = self._get_dict_from_schema()
+
+ keys, dummy_labels = list(zip(*batch))
+ df_batch_labels[self.key_column].extend(keys)
+ df_batch_labels[self.result_columns[0]].extend(dummy_labels)
+
+ return df_batch_labels

DPF/filters/multigpu_filter.py

@@ -61,6 +61,16 @@ def __init__(
  self.devices = devices
  self.num_parts = len(devices)
 
+ # getting result columns names
+ datafilter = self.filter_class(**self.filter_params, device=devices[0]) # type: ignore
+ self._result_columns = datafilter.result_columns
+ del datafilter
+ torch.cuda.empty_cache()
+
+ @property
+ def result_columns(self) -> list[str]:
+ return self._result_columns
+
  def run(
  self,
  df: pd.DataFrame,

DPF/pipelines/__init__.py

@@ -0,0 +1 @@
+from .filter_pipeline import FilterPipeline

DPF/processors/processor.py

@@ -53,6 +53,10 @@ def columns(self) -> list[str]:
  """Columns that presented in dataframe"""
  return self._df.columns.tolist() # type: ignore
 
+ @property
+ def modalities(self) -> list[str]:
+ return list(self.config.modality2datatype.keys())
+
  def __getitem__(self, column_name: str) -> pd.Series:
  return self._df[column_name]
 
@@ -63,7 +67,7 @@ def print_summary(self) -> None:
  """Prints summary info about dataset"""
  print('Dataset format:', config2format(self.config))
  print('Path:', self.config.path)
- print('Modalities:', list(self.config.modality2datatype.keys()))
+ print('Modalities:', self.modalities)
 
  cols = self.columns
  print('Columns:', len(cols))

DPF/utils/logger.py

@@ -2,8 +2,9 @@
 import logging.config
 import os
 import sys
+from typing import Any
 
-LOGGERS_CONFIG = {
+LOGGERS_CONFIG: dict[str, Any] = {
  "version": 1,
  "disable_existing_loggers": False,
  "formatters": {"default": {"format": "[%(asctime)s][%(levelname)s]: %(message)s"}},
@@ -27,17 +28,42 @@
 }
 
 
+def get_logging_config_copy() -> dict[str, Any]:
+ config = LOGGERS_CONFIG.copy()
+ config["loggers"] = LOGGERS_CONFIG["loggers"].copy()
+ config["handlers"] = LOGGERS_CONFIG["handlers"].copy()
+ config["formatters"] = LOGGERS_CONFIG["formatters"].copy()
+ return config
+
+
 def init_logger(filename: str, logger_name: str = "filter_logger", logging_dir: str = "./logs/") -> logging.Logger:
  os.makedirs(logging_dir, exist_ok=True)
 
- LOGGERS_CONFIG["handlers"]["file"]["filename"] = os.path.join(logging_dir, filename) # type: ignore [index]
- LOGGERS_CONFIG["loggers"][logger_name] = { # type: ignore [index]
+ config = get_logging_config_copy()
+
+ config["handlers"]["file"]["filename"] = os.path.join(logging_dir, filename)
+ config["loggers"][logger_name] = {
  "handlers": ["console", "file"],
  "level": logging.DEBUG,
  }
 
- logging.config.dictConfig(LOGGERS_CONFIG)
+ logging.config.dictConfig(config)
+ logger = logging.getLogger(logger_name)
+
+ return logger
+
+
+def init_stdout_logger(logger_name: str = "filter_logger") -> logging.Logger:
+ config = get_logging_config_copy()
+
+ config["loggers"][logger_name] = {
+ "handlers": ["console"],
+ "level": logging.DEBUG,
+ }
+ config["handlers"].pop("file")
+ config["loggers"].pop("template")
+
+ logging.config.dictConfig(config)
  logger = logging.getLogger(logger_name)
- logger.info("Logger initialized")
 
  return logger

tests/test_pipelines.py

@@ -0,0 +1,118 @@
+from DPF import DatasetReader
+from DPF.configs import ShardsDatasetConfig
+from DPF.filters.images.dummy_gpu_filter import DummyGPUFilter
+from DPF.filters.images.hash_filters import PHashFilter
+from DPF.filters.images.info_filter import ImageInfoFilter
+from DPF.pipelines import FilterPipeline
+
+
+def test_pipeline_imageinfo():
+ path = 'tests/datasets/shards_correct'
+ config = ShardsDatasetConfig.from_path_and_columns(
+ path,
+ image_name_col="image_name",
+ text_col="caption"
+ )
+
+ reader = DatasetReader()
+ processor = reader.read_from_config(config)
+
+ pipeline = FilterPipeline("test_pipeline_imageinfo")
+ pipeline.add_datafilter(
+ ImageInfoFilter,
+ {'workers': 1},
+ )
+ pipeline.run(processor)
+
+ assert len({'width', 'height', 'channels'}.intersection(set(processor.df.columns))) == 3
+
+
+def test_pipeline_imageinfo_bad_1():
+ path = 'tests/datasets/shards_bad_image'
+ config = ShardsDatasetConfig.from_path_and_columns(
+ path,
+ image_name_col="image_name",
+ text_col="caption"
+ )
+
+ reader = DatasetReader()
+ processor = reader.read_from_config(config)
+
+ pipeline = FilterPipeline("test_pipeline_imageinfo_bad")
+ pipeline.add_datafilter(
+ ImageInfoFilter,
+ {'workers': 1},
+ on_error="continue"
+ )
+ pipeline.run(processor)
+
+ assert len({'width', 'height', 'channels'}.intersection(set(processor.df.columns))) == 0
+
+
+def test_pipeline_imageinfo_bad_2():
+ path = 'tests/datasets/shards_bad_image'
+ config = ShardsDatasetConfig.from_path_and_columns(
+ path,
+ image_name_col="image_name",
+ text_col="caption"
+ )
+
+ reader = DatasetReader()
+ processor = reader.read_from_config(config)
+
+ pipeline = FilterPipeline("test_pipeline_imageinfo_bad")
+ pipeline.add_datafilter(
+ ImageInfoFilter,
+ {'workers': 1},
+ on_error="stop"
+ )
+ error = None
+ try:
+ pipeline.run(processor)
+ except Exception as err:
+ error = err
+ assert error is not None
+
+
+def test_pipeline_phash_dedup():
+ path = 'tests/datasets/shards_correct'
+ config = ShardsDatasetConfig.from_path_and_columns(
+ path,
+ image_name_col="image_name",
+ text_col="caption"
+ )
+
+ reader = DatasetReader()
+ processor = reader.read_from_config(config)
+
+ pipeline = FilterPipeline("test_pipeline_phash_dedup")
+ pipeline.add_datafilter(
+ PHashFilter,
+ {'workers': 1},
+ )
+ pipeline.add_deduplication(['image_phash_8'])
+ pipeline.run(processor)
+
+ assert len(processor.df) == 1
+
+
+def test_pipeline_multigpu():
+ path = 'tests/datasets/shards_correct'
+ config = ShardsDatasetConfig.from_path_and_columns(
+ path,
+ image_name_col="image_name",
+ text_col="caption"
+ )
+
+ reader = DatasetReader()
+ processor = reader.read_from_config(config)
+
+ pipeline = FilterPipeline("test_pipeline_phash_dedup")
+ pipeline.add_datafilter(
+ DummyGPUFilter,
+ {'workers': 1},
+ devices=["cuda:0", "cuda:1"]
+ )
+ pipeline.run(processor)
+
+ assert len(processor.df) == 2 and 'dummy_label' in processor.columns