Skip to content

Commit

Permalink
Simplify folder 3d dataset (#2184)
Browse files Browse the repository at this point in the history
  • Loading branch information
@samet-akcay
samet-akcay authored Jul 11, 2024
1 parent 4da6605 commit 70c1ef0
Showing 1 changed file with 72 additions and 111 deletions.
183 changes: 72 additions & 111 deletions src/anomalib/data/depth/folder_3d.py
View file
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 = (

Check warning on line 131 in src/anomalib/data/depth/folder_3d.py

View check run for this annotation

Codecov / codecov/patch

src/anomalib/data/depth/folder_3d.py#L131 

Added line #L131 was not covered by tests
"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

0 comments on commit 70c1ef0

Please sign in to comment.