DataProcessingFramework
97 строк · 3.2 Кб
1import os
2from typing import Optional, Union
3
4from DPF.datatypes import ColumnDataType, DataType, FileDataType
5from DPF.modalities import MODALITIES, ModalityName
6
7from .dataset_config import DatasetConfig
8
9
10class FilesDatasetConfig(DatasetConfig):
11"""Config for Files dataset type"""
12
13def __init__(
14self,
15path: str,
16datatypes: list[Union[FileDataType, ColumnDataType]],
17):
18"""
19Parameters
20----------
21path: str
22Path to dataset metadata file
23datatypes: list[Union[FileDataType, ColumnDataType]]
24List of datatypes in dataset
25"""
26super().__init__(path)
27self.table_path = path
28self.base_path = os.path.dirname(self.table_path)
29self._datatypes = datatypes
30self._modality2datatype = {d.modality.name: d for d in datatypes}
31
32assert len({d.modality.name for d in datatypes}) == len(datatypes), \
33"More than one datatype with same modality is not supported"
34for data in self.datatypes:
35assert isinstance(data, (ColumnDataType, FileDataType))
36
37@property
38def datatypes(self) -> list[DataType]:
39return self._datatypes # type: ignore
40
41@property
42def modality2datatype(self) -> dict[ModalityName, DataType]:
43return self._modality2datatype # type: ignore
44
45@property
46def user_column2default_column(self) -> dict[str, str]:
47mapping = {}
48for data in self.datatypes:
49if isinstance(data, ColumnDataType):
50mapping[data.user_column_name] = data.column_name
51elif isinstance(data, FileDataType):
52mapping[data.user_path_column_name] = data.modality.path_column
53return mapping
54
55@classmethod
56def from_path_and_columns(
57cls,
58path: str,
59image_path_col: Optional[str] = None,
60video_path_col: Optional[str] = None,
61text_col: Optional[str] = None,
62) -> "FilesDatasetConfig":
63"""
64Parameters
65----------
66path: str
67Path to dataset metadata file
68image_path_col: Optional[str] = None
69Name of column with image paths
70video_path_col: Optional[str] = None
71Name of column with video paths
72text_col: Optional[str] = None
73Name of column with text
74
75Returns
76-------
77FilesDatasetConfig
78Instance of itself
79"""
80datatypes: list[Union[FileDataType, ColumnDataType]] = []
81if image_path_col:
82datatypes.append(FileDataType(MODALITIES['image'], image_path_col))
83if video_path_col:
84datatypes.append(FileDataType(MODALITIES['video'], video_path_col))
85if text_col:
86datatypes.append(ColumnDataType(MODALITIES['text'], text_col))
87assert len(datatypes) > 0, "At least one modality should be provided"
88return cls(path, datatypes)
89
90def __repr__(self) -> str:
91s = "FilesDatasetConfig(\n\t"
92s += f'table_path="{self.table_path}",\n\t'
93s += 'datatypes=[\n\t\t'
94s += '\n\t\t'.join([str(i) for i in self.datatypes])
95s += '\n\t]'
96s += '\n)'
97return s
98