Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
165 changes: 99 additions & 66 deletions src/anomalib/data/depth/folder_3d.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,62 +25,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 @@ -100,24 +54,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 @@ -152,7 +99,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 @@ -168,6 +114,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