Skip to content

Commit 07a7af2

Browse files
committed
Refactored-make_folder3d_dataset-ruff-error-C901
Signed-off-by: sahusiddharth <[email protected]>
1 parent ca90807 commit 07a7af2

File tree

1 file changed

+99
-66
lines changed

1 file changed

+99
-66
lines changed

src/anomalib/data/depth/folder_3d.py

Lines changed: 99 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -25,62 +25,16 @@
2525
from 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

Comments
 (0)