2525from anomalib .data .utils .path import _prepare_files_labels , validate_and_resolve_path
2626
2727
28- def make_folder3d_dataset ( # noqa: C901
28+ def make_path_dirs (
2929 normal_dir : str | Path ,
30- root : str | Path | None = None ,
3130 abnormal_dir : str | Path | None = None ,
3231 normal_test_dir : str | Path | None = None ,
3332 mask_dir : str | Path | None = None ,
3433 normal_depth_dir : str | Path | None = None ,
3534 abnormal_depth_dir : str | Path | None = None ,
3635 normal_test_depth_dir : str | Path | None = None ,
37- split : str | Split | None = None ,
38- extensions : tuple [str , ...] | None = None ,
39- ) -> DataFrame :
40- """Make Folder Dataset.
41-
42- Args:
43- normal_dir (str | Path): Path to the directory containing normal images.
44- root (str | Path | None): Path to the root directory of the dataset.
45- Defaults to ``None``.
46- abnormal_dir (str | Path | None, optional): Path to the directory containing abnormal images.
47- Defaults to ``None``.
48- normal_test_dir (str | Path | None, optional): Path to the directory containing normal images for the test
49- dataset. Normal test images will be a split of `normal_dir` if `None`.
50- Defaults to ``None``.
51- mask_dir (str | Path | None, optional): Path to the directory containing the mask annotations.
52- Defaults to ``None``.
53- normal_depth_dir (str | Path | None, optional): Path to the directory containing
54- normal depth images for the test dataset. Normal test depth images will be a split of `normal_dir`
55- Defaults to ``None``.
56- abnormal_depth_dir (str | Path | None, optional): Path to the directory containing abnormal depth images for
57- the test dataset.
58- Defaults to ``None``.
59- normal_test_depth_dir (str | Path | None, optional): Path to the directory containing normal depth images for
60- the test dataset. Normal test images will be a split of `normal_dir` if `None`.
61- Defaults to ``None``.
62- split (str | Split | None, optional): Dataset split (ie., Split.FULL, Split.TRAIN or Split.TEST).
63- Defaults to ``None``.
64- extensions (tuple[str, ...] | None, optional): Type of the image extensions to read from the directory.
65- Defaults to ``None``.
66-
67- Returns:
68- DataFrame: an output dataframe containing samples for the requested split (ie., train or test)
69- """
70- normal_dir = validate_and_resolve_path (normal_dir , root )
71- abnormal_dir = validate_and_resolve_path (abnormal_dir , root ) if abnormal_dir else None
72- normal_test_dir = validate_and_resolve_path (normal_test_dir , root ) if normal_test_dir else None
73- mask_dir = validate_and_resolve_path (mask_dir , root ) if mask_dir else None
74- normal_depth_dir = validate_and_resolve_path (normal_depth_dir , root ) if normal_depth_dir else None
75- abnormal_depth_dir = validate_and_resolve_path (abnormal_depth_dir , root ) if abnormal_depth_dir else None
76- normal_test_depth_dir = validate_and_resolve_path (normal_test_depth_dir , root ) if normal_test_depth_dir else None
77-
78- if not normal_dir .is_dir ():
79- msg = "A folder location must be provided in normal_dir."
80- raise ValueError (msg )
81-
82- filenames = []
83- labels = []
36+ ) -> dict :
37+ """Create a dictionary containing paths to different directories."""
8438 dirs = {DirType .NORMAL : normal_dir }
8539
8640 if abnormal_dir :
@@ -100,24 +54,17 @@ def make_folder3d_dataset( # noqa: C901
10054
10155 if mask_dir :
10256 dirs [DirType .MASK ] = mask_dir
57+ return dirs
10358
104- for dir_type , path in dirs .items ():
105- filename , label = _prepare_files_labels (path , dir_type , extensions )
106- filenames += filename
107- labels += label
108-
109- samples = DataFrame ({"image_path" : filenames , "label" : labels })
110- samples = samples .sort_values (by = "image_path" , ignore_index = True )
111-
112- # Create label index for normal (0) and abnormal (1) images.
113- samples .loc [
114- (samples .label == DirType .NORMAL ) | (samples .label == DirType .NORMAL_TEST ),
115- "label_index" ,
116- ] = LabelName .NORMAL
117- samples .loc [(samples .label == DirType .ABNORMAL ), "label_index" ] = LabelName .ABNORMAL
118- samples .label_index = samples .label_index .astype ("Int64" )
11959
120- # If a path to mask is provided, add it to the sample dataframe.
60+ def add_mask (
61+ samples : DataFrame ,
62+ normal_depth_dir : str | Path | None = None ,
63+ abnormal_dir : str | Path | None = None ,
64+ normal_test_dir : str | Path | None = None ,
65+ mask_dir : str | Path | None = None ,
66+ ) -> DataFrame :
67+ """If a path to mask is provided, add it to the sample dataframe."""
12168 if normal_depth_dir :
12269 samples .loc [samples .label == DirType .NORMAL , "depth_path" ] = samples .loc [
12370 samples .label == DirType .NORMAL_DEPTH
@@ -152,7 +99,6 @@ def make_folder3d_dataset( # noqa: C901
15299
153100 samples = samples .astype ({"depth_path" : "str" })
154101
155- # If a path to mask is provided, add it to the sample dataframe.
156102 if mask_dir and abnormal_dir :
157103 samples .loc [samples .label == DirType .ABNORMAL , "mask_path" ] = samples .loc [
158104 samples .label == DirType .MASK
@@ -168,6 +114,93 @@ def make_folder3d_dataset( # noqa: C901
168114 raise FileNotFoundError (msg )
169115 else :
170116 samples ["mask_path" ] = ""
117+ return samples
118+
119+
120+ def make_folder3d_dataset (
121+ normal_dir : str | Path ,
122+ root : str | Path | None = None ,
123+ abnormal_dir : str | Path | None = None ,
124+ normal_test_dir : str | Path | None = None ,
125+ mask_dir : str | Path | None = None ,
126+ normal_depth_dir : str | Path | None = None ,
127+ abnormal_depth_dir : str | Path | None = None ,
128+ normal_test_depth_dir : str | Path | None = None ,
129+ split : str | Split | None = None ,
130+ extensions : tuple [str , ...] | None = None ,
131+ ) -> DataFrame :
132+ """Make Folder Dataset.
133+
134+ Args:
135+ normal_dir (str | Path): Path to the directory containing normal images.
136+ root (str | Path | None): Path to the root directory of the dataset.
137+ Defaults to ``None``.
138+ abnormal_dir (str | Path | None, optional): Path to the directory containing abnormal images.
139+ Defaults to ``None``.
140+ normal_test_dir (str | Path | None, optional): Path to the directory containing normal images for the test
141+ dataset. Normal test images will be a split of `normal_dir` if `None`.
142+ Defaults to ``None``.
143+ mask_dir (str | Path | None, optional): Path to the directory containing the mask annotations.
144+ Defaults to ``None``.
145+ normal_depth_dir (str | Path | None, optional): Path to the directory containing
146+ normal depth images for the test dataset. Normal test depth images will be a split of `normal_dir`
147+ Defaults to ``None``.
148+ abnormal_depth_dir (str | Path | None, optional): Path to the directory containing abnormal depth images for
149+ the test dataset.
150+ Defaults to ``None``.
151+ normal_test_depth_dir (str | Path | None, optional): Path to the directory containing normal depth images for
152+ the test dataset. Normal test images will be a split of `normal_dir` if `None`.
153+ Defaults to ``None``.
154+ split (str | Split | None, optional): Dataset split (ie., Split.FULL, Split.TRAIN or Split.TEST).
155+ Defaults to ``None``.
156+ extensions (tuple[str, ...] | None, optional): Type of the image extensions to read from the directory.
157+ Defaults to ``None``.
158+
159+ Returns:
160+ DataFrame: an output dataframe containing samples for the requested split (ie., train or test)
161+ """
162+ normal_dir = validate_and_resolve_path (normal_dir , root )
163+ abnormal_dir = validate_and_resolve_path (abnormal_dir , root ) if abnormal_dir else None
164+ normal_test_dir = validate_and_resolve_path (normal_test_dir , root ) if normal_test_dir else None
165+ mask_dir = validate_and_resolve_path (mask_dir , root ) if mask_dir else None
166+ normal_depth_dir = validate_and_resolve_path (normal_depth_dir , root ) if normal_depth_dir else None
167+ abnormal_depth_dir = validate_and_resolve_path (abnormal_depth_dir , root ) if abnormal_depth_dir else None
168+ normal_test_depth_dir = validate_and_resolve_path (normal_test_depth_dir , root ) if normal_test_depth_dir else None
169+
170+ if not normal_dir .is_dir ():
171+ msg = "A folder location must be provided in normal_dir."
172+ raise ValueError (msg )
173+
174+ filenames = []
175+ labels = []
176+ dirs = make_path_dirs (
177+ normal_dir ,
178+ abnormal_dir ,
179+ normal_test_dir ,
180+ mask_dir ,
181+ normal_depth_dir ,
182+ abnormal_depth_dir ,
183+ normal_test_depth_dir ,
184+ )
185+
186+ for dir_type , path in dirs .items ():
187+ filename , label = _prepare_files_labels (path , dir_type , extensions )
188+ filenames += filename
189+ labels += label
190+
191+ samples = DataFrame ({"image_path" : filenames , "label" : labels })
192+ samples = samples .sort_values (by = "image_path" , ignore_index = True )
193+
194+ # Create label index for normal (0) and abnormal (1) images.
195+ samples .loc [
196+ (samples .label == DirType .NORMAL ) | (samples .label == DirType .NORMAL_TEST ),
197+ "label_index" ,
198+ ] = LabelName .NORMAL
199+ samples .loc [(samples .label == DirType .ABNORMAL ), "label_index" ] = LabelName .ABNORMAL
200+ samples .label_index = samples .label_index .astype ("Int64" )
201+
202+ # If a path to mask is provided, add it to the sample dataframe.
203+ samples = add_mask (samples , normal_depth_dir , abnormal_dir , normal_test_dir , mask_dir )
171204
172205 # remove all the rows with temporal image samples that have already been assigned
173206 samples = samples .loc [
0 commit comments