Skip to content

Commit 4da6605

Browse files
Refactored-make_folder3d_dataset-ruff-error-C901 (#1926)
Signed-off-by: sahusiddharth <[email protected]>
1 parent d094d4b commit 4da6605

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
@@ -24,62 +24,16 @@
2424
from 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

Comments
 (0)