feat: add transforms in pipelines

Author: boomb0om

DPF/pipelines/filter_pipeline.py

@@ -6,6 +6,7 @@
 from DPF.filters import ColumnFilter, DataFilter
 from DPF.filters.multigpu_filter import MultiGPUDataFilter
 from DPF.processors import DatasetProcessor
+from DPF.transforms import BaseFilesTransforms
 from DPF.utils.logger import init_logger, init_stdout_logger
 
 from .pipeline_stages import (
@@ -14,6 +15,7 @@
  FilterPipelineStage,
  PipelineStage,
  ShufflePipelineStage,
+ TransformPipelineStage,
 )
 from .types import OnErrorOptions
 
@@ -39,25 +41,25 @@ def add_datafilter(
  datafilter: type[DataFilter],
  datafilter_kwargs: dict[str, Any],
  devices: Optional[list[str]] = None,
- processor_run_kwargs: Optional[dict[str, Any]] = None,
+ processor_apply_kwargs: Optional[dict[str, Any]] = None,
  on_error: OnErrorOptions = "stop",
  skip_if_columns_exist: bool = True
  ) -> None:
- if processor_run_kwargs is None:
- processor_run_kwargs = {}
+ if processor_apply_kwargs is None:
+ processor_apply_kwargs = {}
 
  if devices is None:
  stage = FilterPipelineStage(
  'datafilter', filter_class=datafilter,
- filter_kwargs=datafilter_kwargs, processor_run_kwargs=processor_run_kwargs,
+ filter_kwargs=datafilter_kwargs, processor_apply_kwargs=processor_apply_kwargs,
  skip_if_columns_exist=skip_if_columns_exist
  )
  elif len(devices) == 0:
  new_kwargs = datafilter_kwargs.copy()
  new_kwargs['device'] = devices[0]
  stage = FilterPipelineStage(
  'datafilter', filter_class=datafilter,
- filter_kwargs=new_kwargs, processor_run_kwargs=processor_run_kwargs,
+ filter_kwargs=new_kwargs, processor_apply_kwargs=processor_apply_kwargs,
  skip_if_columns_exist=skip_if_columns_exist
  )
  else:
@@ -68,7 +70,7 @@ def add_datafilter(
  "datafilter_class": datafilter,
  "datafilter_params": datafilter_kwargs
  },
- processor_run_kwargs=processor_run_kwargs,
+ processor_apply_kwargs=processor_apply_kwargs,
  skip_if_columns_exist=skip_if_columns_exist
  )
 
@@ -80,16 +82,16 @@ def add_columnfilter(
  self,
  columnfilter: type[ColumnFilter],
  columnfilter_kwargs: dict[str, Any],
- processor_run_kwargs: Optional[dict[str, Any]] = None,
+ processor_apply_kwargs: Optional[dict[str, Any]] = None,
  on_error: OnErrorOptions = "stop",
  skip_if_columns_exist: bool = True
  ) -> None:
- if processor_run_kwargs is None:
- processor_run_kwargs = {}
+ if processor_apply_kwargs is None:
+ processor_apply_kwargs = {}
 
  stage = FilterPipelineStage(
  'columnfilter', filter_class=columnfilter,
- filter_kwargs=columnfilter_kwargs, processor_run_kwargs=processor_run_kwargs,
+ filter_kwargs=columnfilter_kwargs, processor_apply_kwargs=processor_apply_kwargs,
  skip_if_columns_exist=skip_if_columns_exist
  )
 
@@ -123,6 +125,21 @@ def add_dataframe_filter(
  PipelineStageRunner(stage, on_error=on_error)
  )
 
+ def add_transforms(
+ self,
+ transforms_class: type[BaseFilesTransforms],
+ transforms_kwargs: dict[str, Any],
+ processor_apply_kwargs: Optional[dict[str, Any]] = None,
+ on_error: OnErrorOptions = "stop"
+ ) -> None:
+ stage = TransformPipelineStage(
+ transforms_class, transforms_kwargs,
+ processor_apply_kwargs=processor_apply_kwargs
+ )
+ self.stages.append(
+ PipelineStageRunner(stage, on_error=on_error)
+ )
+
  def _log_dataset_info(self, processor: DatasetProcessor) -> None:
  self.logger.info(f'Dataset path: {processor.config.path}')
  self.logger.info(f'Dataset modalities: {processor.modalities}')

DPF/pipelines/pipeline_stages.py

@@ -1,12 +1,13 @@
 import logging
 from abc import ABC, abstractmethod
-from typing import Any, Callable, Union
+from typing import Any, Callable, Optional, Union
 
 import pandas as pd
 
 from DPF.filters import ColumnFilter, DataFilter
 from DPF.filters.multigpu_filter import MultiGPUDataFilter
 from DPF.processors import DatasetProcessor
+from DPF.transforms import BaseFilesTransforms
 
 from .types import FilterTypes
 
@@ -66,13 +67,17 @@ def __init__(
  filter_type: FilterTypes,
  filter_class: Union[type[DataFilter], type[ColumnFilter], type[MultiGPUDataFilter]],
  filter_kwargs: dict[str, Any],
- processor_run_kwargs: dict[str, Any],
+ processor_apply_kwargs: Optional[dict[str, Any]] = None,
  skip_if_columns_exist: bool = True
  ):
  self.filter_type = filter_type
  self.filter_class = filter_class
  self.filter_kwargs = filter_kwargs
- self.processor_run_kwargs = processor_run_kwargs
+
+ self.processor_apply_kwargs = processor_apply_kwargs
+ if self.processor_apply_kwargs is None:
+ self.processor_apply_kwargs = {}
+
  self.skip_if_columns_exist = skip_if_columns_exist
 
  @property
@@ -96,10 +101,35 @@ def run(self, processor: DatasetProcessor, logger: logging.Logger) -> None:
  processor.df.drop(columns=columns_to_be_added, inplace=True, errors='ignore')
 
  if self.filter_type == 'datafilter':
- processor.apply_data_filter(filter_obj, **self.processor_run_kwargs) # type: ignore
+ processor.apply_data_filter(filter_obj, **self.processor_apply_kwargs) # type: ignore
  elif self.filter_type == 'columnfilter':
- processor.apply_column_filter(filter_obj, **self.processor_run_kwargs) # type: ignore
+ processor.apply_column_filter(filter_obj, **self.processor_apply_kwargs) # type: ignore
  elif self.filter_type == 'multigpufilter':
- processor.apply_multi_gpu_data_filter(filter_obj, **self.processor_run_kwargs)
+ processor.apply_multi_gpu_data_filter(filter_obj, **self.processor_apply_kwargs) # type: ignore
  else:
  raise ValueError(f"Unknown filter type: {self.filter_type}")
+
+
+class TransformPipelineStage(PipelineStage):
+
+ def __init__(
+ self,
+ transforms_class: type[BaseFilesTransforms],
+ transforms_kwargs: dict[str, Any],
+ processor_apply_kwargs: Optional[dict[str, Any]] = None,
+ ):
+ self.transforms_class = transforms_class
+ self.transforms_kwargs = transforms_kwargs
+
+ self.processor_apply_kwargs = processor_apply_kwargs
+ if self.processor_apply_kwargs is None:
+ self.processor_apply_kwargs = {}
+
+ @property
+ def stage_name(self) -> str:
+ return f"TransformPipelineStage(transforms_class={self.transforms_class}, transforms_kwargs={self.transforms_kwargs})"
+
+ def run(self, processor: DatasetProcessor, logger: logging.Logger) -> None:
+ transforms = self.transforms_class(**self.transforms_kwargs)
+
+ processor.apply_transform(transforms, **self.processor_apply_kwargs) # type: ignore