Skip to content

Simplify folder 3d dataset #2184

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
samet-akcay merged 1 commit into from
Jul 11, 2024
Merged

Simplify folder 3d dataset #2184

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
183 changes: 72 additions & 111 deletions src/anomalib/data/depth/folder_3d.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,98 +24,6 @@
from anomalib.data.utils.path import _prepare_files_labels, validate_and_resolve_path


def make_path_dirs(
normal_dir: str | Path,
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,
) -> dict:
"""Create a dictionary containing paths to different directories."""
dirs = {DirType.NORMAL: normal_dir}

if abnormal_dir:
dirs[DirType.ABNORMAL] = abnormal_dir

if normal_test_dir:
dirs[DirType.NORMAL_TEST] = normal_test_dir

if normal_depth_dir:
dirs[DirType.NORMAL_DEPTH] = normal_depth_dir

if abnormal_depth_dir:
dirs[DirType.ABNORMAL_DEPTH] = abnormal_depth_dir

if normal_test_depth_dir:
dirs[DirType.NORMAL_TEST_DEPTH] = normal_test_depth_dir

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


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
].image_path.to_numpy()
samples.loc[samples.label == DirType.ABNORMAL, "depth_path"] = samples.loc[
samples.label == DirType.ABNORMAL_DEPTH
].image_path.to_numpy()

if normal_test_dir:
samples.loc[samples.label == DirType.NORMAL_TEST, "depth_path"] = samples.loc[
samples.label == DirType.NORMAL_TEST_DEPTH
].image_path.to_numpy()

# make sure every rgb image has a corresponding depth image and that the file exists
mismatch = (
samples.loc[samples.label_index == LabelName.ABNORMAL]
.apply(lambda x: Path(x.image_path).stem in Path(x.depth_path).stem, axis=1)
.all()
)
if not mismatch:
msg = """Mismatch between anomalous images and depth images. Make sure the mask files
in 'xyz' folder follow the same naming convention as the anomalous images in the dataset
(e.g. image: '000.png', depth: '000.tiff')."""
raise MisMatchError(msg)

missing_depth_files = samples.depth_path.apply(
lambda x: Path(x).exists() if not isna(x) else True,
).all()
if not missing_depth_files:
msg = "Missing depth image files."
raise FileNotFoundError(msg)

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

if mask_dir and abnormal_dir:
samples.loc[samples.label == DirType.ABNORMAL, "mask_path"] = samples.loc[
samples.label == DirType.MASK
].image_path.to_numpy()
samples["mask_path"] = samples["mask_path"].fillna("")
samples = samples.astype({"mask_path": "str"})

# make sure all the files exist
if not samples.mask_path.apply(
lambda x: Path(x).exists() if x != "" else True,
).all():
msg = f"Missing mask files. mask_dir={mask_dir}"
raise FileNotFoundError(msg)
else:
samples["mask_path"] = ""
return samples


def make_folder3d_dataset(
normal_dir: str | Path,
root: str | Path | None = None,
Expand Down Expand Up @@ -170,27 +78,28 @@ def make_folder3d_dataset(
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
dirs = {
DirType.NORMAL: normal_dir,
DirType.ABNORMAL: abnormal_dir,
DirType.NORMAL_TEST: normal_test_dir,
DirType.NORMAL_DEPTH: normal_depth_dir,
DirType.ABNORMAL_DEPTH: abnormal_depth_dir,
DirType.NORMAL_TEST_DEPTH: normal_test_depth_dir,
DirType.MASK: mask_dir,
}

filenames: list[Path] = []
labels: list[str] = []

for dir_type, dir_path in dirs.items():
if dir_path is not None:
filename, label = _prepare_files_labels(dir_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",
Expand All @@ -199,9 +108,61 @@ def make_folder3d_dataset(
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)
if normal_depth_dir:
samples.loc[samples.label == DirType.NORMAL, "depth_path"] = samples.loc[
samples.label == DirType.NORMAL_DEPTH
].image_path.to_numpy()
samples.loc[samples.label == DirType.ABNORMAL, "depth_path"] = samples.loc[
samples.label == DirType.ABNORMAL_DEPTH
].image_path.to_numpy()

if normal_test_dir:
samples.loc[samples.label == DirType.NORMAL_TEST, "depth_path"] = samples.loc[
samples.label == DirType.NORMAL_TEST_DEPTH
].image_path.to_numpy()

# make sure every rgb image has a corresponding depth image and that the file exists
mismatch = (
samples.loc[samples.label_index == LabelName.ABNORMAL]
.apply(lambda x: Path(x.image_path).stem in Path(x.depth_path).stem, axis=1)
.all()
)
if not mismatch:
msg = (
"Mismatch between anomalous images and depth images. "
"Make sure the mask files in 'xyz' folder follow the same naming "
"convention as the anomalous images in the dataset"
"(e.g. image: '000.png', depth: '000.tiff')."
)
raise MisMatchError(msg)

missing_depth_files = samples.depth_path.apply(
lambda x: Path(x).exists() if not isna(x) else True,
).all()
if not missing_depth_files:
msg = "Missing depth image files."
raise FileNotFoundError(msg)

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
].image_path.to_numpy()
samples["mask_path"] = samples["mask_path"].fillna("")
samples = samples.astype({"mask_path": "str"})

# Make sure all the files exist
if not samples.mask_path.apply(
lambda x: Path(x).exists() if x != "" else True,
).all():
msg = f"Missing mask files. mask_dir={mask_dir}"
raise FileNotFoundError(msg)
else:
samples["mask_path"] = ""

# remove all the rows with temporal image samples that have already been assigned
# Remove all the rows with temporal image samples that have already been assigned
samples = samples.loc[
(samples.label == DirType.NORMAL) | (samples.label == DirType.ABNORMAL) | (samples.label == DirType.NORMAL_TEST)
]
Expand Down