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