Skip to content

Commit

Permalink
Refactored-make_folder3d_dataset-ruff-error-C901 (#1926)
Browse files Browse the repository at this point in the history 
Signed-off-by: sahusiddharth <[email protected]>
  • Loading branch information
@sahusiddharth
sahusiddharth authored Jul 11, 2024
1 parent d094d4b commit 4da6605
Showing 1 changed file with 99 additions and 66 deletions.
165 changes: 99 additions & 66 deletions src/anomalib/data/depth/folder_3d.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,62 +24,16 @@
from anomalib.data.utils.path import _prepare_files_labels, validate_and_resolve_path


def make_folder3d_dataset( # noqa: C901
def make_path_dirs(
normal_dir: str | Path,
root: str | Path | None = None,
abnormal_dir: str | Path | None = None,
normal_test_dir: str | Path | None = None,
mask_dir: str | Path | None = None,
normal_depth_dir: str | Path | None = None,
abnormal_depth_dir: str | Path | None = None,
normal_test_depth_dir: str | Path | None = None,
split: str | Split | None = None,
extensions: tuple[str, ...] | None = None,
) -> DataFrame:
"""Make Folder Dataset.
Args:
normal_dir (str | Path): Path to the directory containing normal images.
root (str | Path | None): Path to the root directory of the dataset.
Defaults to ``None``.
abnormal_dir (str | Path | None, optional): Path to the directory containing abnormal images.
Defaults to ``None``.
normal_test_dir (str | Path | None, optional): Path to the directory containing normal images for the test
dataset. Normal test images will be a split of `normal_dir` if `None`.
Defaults to ``None``.
mask_dir (str | Path | None, optional): Path to the directory containing the mask annotations.
Defaults to ``None``.
normal_depth_dir (str | Path | None, optional): Path to the directory containing
normal depth images for the test dataset. Normal test depth images will be a split of `normal_dir`
Defaults to ``None``.
abnormal_depth_dir (str | Path | None, optional): Path to the directory containing abnormal depth images for
the test dataset.
Defaults to ``None``.
normal_test_depth_dir (str | Path | None, optional): Path to the directory containing normal depth images for
the test dataset. Normal test images will be a split of `normal_dir` if `None`.
Defaults to ``None``.
split (str | Split | None, optional): Dataset split (ie., Split.FULL, Split.TRAIN or Split.TEST).
Defaults to ``None``.
extensions (tuple[str, ...] | None, optional): Type of the image extensions to read from the directory.
Defaults to ``None``.
Returns:
DataFrame: an output dataframe containing samples for the requested split (ie., train or test)
"""
normal_dir = validate_and_resolve_path(normal_dir, root)
abnormal_dir = validate_and_resolve_path(abnormal_dir, root) if abnormal_dir else None
normal_test_dir = validate_and_resolve_path(normal_test_dir, root) if normal_test_dir else None
mask_dir = validate_and_resolve_path(mask_dir, root) if mask_dir else None
normal_depth_dir = validate_and_resolve_path(normal_depth_dir, root) if normal_depth_dir else None
abnormal_depth_dir = validate_and_resolve_path(abnormal_depth_dir, root) if abnormal_depth_dir else None
normal_test_depth_dir = validate_and_resolve_path(normal_test_depth_dir, root) if normal_test_depth_dir else None

if not normal_dir.is_dir():
msg = "A folder location must be provided in normal_dir."
raise ValueError(msg)

filenames = []
labels = []
) -> dict:
"""Create a dictionary containing paths to different directories."""
dirs = {DirType.NORMAL: normal_dir}

if abnormal_dir:
Expand All @@ -99,24 +53,17 @@ def make_folder3d_dataset( # noqa: C901

if mask_dir:
dirs[DirType.MASK] = mask_dir
return dirs

for dir_type, path in dirs.items():
filename, label = _prepare_files_labels(path, dir_type, extensions)
filenames += filename
labels += label

samples = DataFrame({"image_path": filenames, "label": labels})
samples = samples.sort_values(by="image_path", ignore_index=True)

# Create label index for normal (0) and abnormal (1) images.
samples.loc[
(samples.label == DirType.NORMAL) | (samples.label == DirType.NORMAL_TEST),
"label_index",
] = LabelName.NORMAL
samples.loc[(samples.label == DirType.ABNORMAL), "label_index"] = LabelName.ABNORMAL
samples.label_index = samples.label_index.astype("Int64")

# If a path to mask is provided, add it to the sample dataframe.
def add_mask(
samples: DataFrame,
normal_depth_dir: str | Path | None = None,
abnormal_dir: str | Path | None = None,
normal_test_dir: str | Path | None = None,
mask_dir: str | Path | None = None,
) -> DataFrame:
"""If a path to mask is provided, add it to the sample dataframe."""
if normal_depth_dir:
samples.loc[samples.label == DirType.NORMAL, "depth_path"] = samples.loc[
samples.label == DirType.NORMAL_DEPTH
Expand Down Expand Up @@ -151,7 +98,6 @@ def make_folder3d_dataset( # noqa: C901

samples = samples.astype({"depth_path": "str"})

# If a path to mask is provided, add it to the sample dataframe.
if mask_dir and abnormal_dir:
samples.loc[samples.label == DirType.ABNORMAL, "mask_path"] = samples.loc[
samples.label == DirType.MASK
Expand All @@ -167,6 +113,93 @@ def make_folder3d_dataset( # noqa: C901
raise FileNotFoundError(msg)
else:
samples["mask_path"] = ""
return samples


def make_folder3d_dataset(
normal_dir: str | Path,
root: str | Path | None = None,
abnormal_dir: str | Path | None = None,
normal_test_dir: str | Path | None = None,
mask_dir: str | Path | None = None,
normal_depth_dir: str | Path | None = None,
abnormal_depth_dir: str | Path | None = None,
normal_test_depth_dir: str | Path | None = None,
split: str | Split | None = None,
extensions: tuple[str, ...] | None = None,
) -> DataFrame:
"""Make Folder Dataset.
Args:
normal_dir (str | Path): Path to the directory containing normal images.
root (str | Path | None): Path to the root directory of the dataset.
Defaults to ``None``.
abnormal_dir (str | Path | None, optional): Path to the directory containing abnormal images.
Defaults to ``None``.
normal_test_dir (str | Path | None, optional): Path to the directory containing normal images for the test
dataset. Normal test images will be a split of `normal_dir` if `None`.
Defaults to ``None``.
mask_dir (str | Path | None, optional): Path to the directory containing the mask annotations.
Defaults to ``None``.
normal_depth_dir (str | Path | None, optional): Path to the directory containing
normal depth images for the test dataset. Normal test depth images will be a split of `normal_dir`
Defaults to ``None``.
abnormal_depth_dir (str | Path | None, optional): Path to the directory containing abnormal depth images for
the test dataset.
Defaults to ``None``.
normal_test_depth_dir (str | Path | None, optional): Path to the directory containing normal depth images for
the test dataset. Normal test images will be a split of `normal_dir` if `None`.
Defaults to ``None``.
split (str | Split | None, optional): Dataset split (ie., Split.FULL, Split.TRAIN or Split.TEST).
Defaults to ``None``.
extensions (tuple[str, ...] | None, optional): Type of the image extensions to read from the directory.
Defaults to ``None``.
Returns:
DataFrame: an output dataframe containing samples for the requested split (ie., train or test)
"""
normal_dir = validate_and_resolve_path(normal_dir, root)
abnormal_dir = validate_and_resolve_path(abnormal_dir, root) if abnormal_dir else None
normal_test_dir = validate_and_resolve_path(normal_test_dir, root) if normal_test_dir else None
mask_dir = validate_and_resolve_path(mask_dir, root) if mask_dir else None
normal_depth_dir = validate_and_resolve_path(normal_depth_dir, root) if normal_depth_dir else None
abnormal_depth_dir = validate_and_resolve_path(abnormal_depth_dir, root) if abnormal_depth_dir else None
normal_test_depth_dir = validate_and_resolve_path(normal_test_depth_dir, root) if normal_test_depth_dir else None

if not normal_dir.is_dir():
msg = "A folder location must be provided in normal_dir."
raise ValueError(msg)

filenames = []
labels = []
dirs = make_path_dirs(
normal_dir,
abnormal_dir,
normal_test_dir,
mask_dir,
normal_depth_dir,
abnormal_depth_dir,
normal_test_depth_dir,
)

for dir_type, path in dirs.items():
filename, label = _prepare_files_labels(path, dir_type, extensions)
filenames += filename
labels += label

samples = DataFrame({"image_path": filenames, "label": labels})
samples = samples.sort_values(by="image_path", ignore_index=True)

# Create label index for normal (0) and abnormal (1) images.
samples.loc[
(samples.label == DirType.NORMAL) | (samples.label == DirType.NORMAL_TEST),
"label_index",
] = LabelName.NORMAL
samples.loc[(samples.label == DirType.ABNORMAL), "label_index"] = LabelName.ABNORMAL
samples.label_index = samples.label_index.astype("Int64")

# If a path to mask is provided, add it to the sample dataframe.
samples = add_mask(samples, normal_depth_dir, abnormal_dir, normal_test_dir, mask_dir)

# remove all the rows with temporal image samples that have already been assigned
samples = samples.loc[
Expand Down

0 comments on commit 4da6605

Please sign in to comment.