DataProcessingFramework

1
from typing import Optional, Union
2

3
from DPF.datatypes import ColumnDataType, ShardedDataType
4
from DPF.modalities import MODALITIES
5

6
from .sharded_config import ShardedDatasetConfig
7

8

9
class ShardedFilesDatasetConfig(ShardedDatasetConfig):
10
    """Config for ShardedFiles dataset type"""
11

12
    def __init__(
13
        self,
14
        path: str,
15
        datatypes: list[Union[ShardedDataType, ColumnDataType]],
16
        datafiles_ext: str = "csv",
17
    ):
18
        """
19
        Parameters
20
        ----------
21
        path: str
22
            Path to directory with shards
23
        datatypes: list[Union[ShardedDataType, ColumnDataType]]
24
            List of datatypes in dataset
25
        datafiles_ext: str = "csv"
26
            Extension of files with metadata in shards
27
        """
28
        super().__init__(path, datatypes, datafiles_ext)
29

30
    @classmethod
31
    def from_path_and_columns(
32
        cls,
33
        path: str,
34
        image_name_col: Optional[str] = None,
35
        video_name_col: Optional[str] = None,
36
        text_col: Optional[str] = None,
37
        datafiles_ext: str = "csv",
38
    ) -> "ShardedFilesDatasetConfig":
39
        """
40
        Parameters
41
        ----------
42
        path: str
43
            Path to directory with shards
44
        image_name_col: Optional[str] = None
45
            Name of column with image filenames in shard
46
        video_name_col: Optional[str] = None
47
            Name of column with video filenames in shard
48
        text_col: Optional[str] = None
49
            Name of column with text
50
        datafiles_ext: str = "csv"
51
            Extension of files with metadata in shards
52

53
        Returns
54
        -------
55
        ShardedFilesDatasetConfig
56
            Instance of itself
57
        """
58
        datatypes: list[Union[ShardedDataType, ColumnDataType]] = []
59
        if image_name_col:
60
            datatypes.append(ShardedDataType(MODALITIES['image'], image_name_col))
61
        if video_name_col:
62
            datatypes.append(ShardedDataType(MODALITIES['video'], video_name_col))
63
        if text_col:
64
            datatypes.append(ColumnDataType(MODALITIES['text'], text_col))
65
        assert len(datatypes) > 0, "At least one modality should be provided"
66
        return cls(path, datatypes, datafiles_ext=datafiles_ext)
67