Skip to content

Commit

Permalink
undo changes
Browse files Browse the repository at this point in the history 
Signed-off-by: Ashwin Vaidya <[email protected]>
  • Loading branch information
@ashwinvaidya17
ashwinvaidya17 committed Oct 7, 2024
1 parent 5e6d8ac commit 7ced8dd
Showing 1 changed file with 14 additions and 56 deletions.
70 changes: 14 additions & 56 deletions src/anomalib/utils/visualization/image.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -170,10 +170,7 @@ def _visualize_batch(self, batch: dict) -> Iterator[GeneratorResult]:
box_labels=batch["box_labels"][i].cpu().numpy() if "box_labels" in batch else None,
normalize=self.normalize,
)
yield GeneratorResult(
image=self.visualize_image(image_result),
file_name=file_name,
)
yield GeneratorResult(image=self.visualize_image(image_result), file_name=file_name)

def visualize_image(self, image_result: ImageResult) -> np.ndarray:
"""Generate the visualization for an image.
Expand Down Expand Up @@ -212,28 +209,17 @@ def _visualize_full(self, image_result: ImageResult) -> np.ndarray:

image_grid.add_image(image_result.image, "Image")
if image_result.gt_boxes is not None:
gt_image = draw_boxes(
np.copy(image_result.image),
image_result.gt_boxes,
color=(255, 0, 0),
)
gt_image = draw_boxes(np.copy(image_result.image), image_result.gt_boxes, color=(255, 0, 0))
image_grid.add_image(image=gt_image, color_map="gray", title="Ground Truth")
image_grid.add_image(
image=gt_image,
color_map="gray",
title="Ground Truth",
)
else:
image_grid.add_image(image_result.image, "Image")
pred_image = draw_boxes(
np.copy(image_result.image),
image_result.normal_boxes,
color=(0, 255, 0),
)
pred_image = draw_boxes(
pred_image,
image_result.anomalous_boxes,
color=(255, 0, 0),
)
pred_image = draw_boxes(np.copy(image_result.image), image_result.normal_boxes, color=(0, 255, 0))
pred_image = draw_boxes(pred_image, image_result.anomalous_boxes, color=(255, 0, 0))
image_grid.add_image(pred_image, "Predictions")
if self.task == TaskType.SEGMENTATION:
if image_result.pred_mask is None:
Expand All @@ -248,11 +234,7 @@ def _visualize_full(self, image_result: ImageResult) -> np.ndarray:
title="Ground Truth",
)
image_grid.add_image(image_result.heat_map, "Predicted Heat Map")
image_grid.add_image(
image=image_result.pred_mask,
color_map="gray",
title="Predicted Mask",
)
image_grid.add_image(image=image_result.gt_mask, color_map="gray", title="Ground Truth")
image_grid.add_image(
image=image_result.segmentations,
title="Segmentation Result",
Expand All @@ -262,15 +244,9 @@ def _visualize_full(self, image_result: ImageResult) -> np.ndarray:
if image_result.heat_map is not None:
image_grid.add_image(image_result.heat_map, "Predicted Heat Map")
if image_result.pred_label:
image_classified = add_anomalous_label(
image_result.image,
image_result.pred_score,
)
image_classified = add_anomalous_label(image_result.image, image_result.pred_score)
else:
image_classified = add_normal_label(
image_result.image,
1 - image_result.pred_score,
)
image_classified = add_normal_label(image_result.image, 1 - image_result.pred_score)
image_grid.add_image(image=image_classified, title="Prediction")
elif self.task == TaskType.EXPLANATION:
image_classified = add_normal_label(
Expand Down Expand Up @@ -303,11 +279,7 @@ def _visualize_simple(self, image_result: ImageResult) -> np.ndarray:
color=(0, 0, 255),
)
if image_result.gt_boxes is not None:
image_with_boxes = draw_boxes(
image=image_with_boxes,
boxes=image_result.gt_boxes,
color=(255, 0, 0),
)
image_with_boxes = draw_boxes(image=image_with_boxes, boxes=image_result.gt_boxes, color=(255, 0, 0))
return image_with_boxes
if self.task == TaskType.SEGMENTATION:
visualization = mark_boundaries(
Expand All @@ -319,15 +291,9 @@ def _visualize_simple(self, image_result: ImageResult) -> np.ndarray:
return (visualization * 255).astype(np.uint8)
if self.task == TaskType.CLASSIFICATION:
if image_result.pred_label:
image_classified = add_anomalous_label(
image_result.image,
image_result.pred_score,
)
image_classified = add_anomalous_label(image_result.image, image_result.pred_score)
else:
image_classified = add_normal_label(
image_result.image,
1 - image_result.pred_score,
)
image_classified = add_normal_label(image_result.image, 1 - image_result.pred_score)
return image_classified
msg = f"Unknown task type: {self.task}"
raise ValueError(msg)
Expand All @@ -345,24 +311,15 @@ def __init__(self) -> None:
self.figure: matplotlib.figure.Figure | None = None
self.axis: Axes | np.ndarray | None = None

def add_image(
self,
image: np.ndarray,
title: str | None = None,
color_map: str | None = None,
) -> None:
def add_image(self, image: np.ndarray, title: str | None = None, color_map: str | None = None) -> None:
"""Add an image to the grid.
Args:
image (np.ndarray): Image which should be added to the figure.
title (str): Image title shown on the plot.
color_map (str | None): Name of matplotlib color map used to map scalar data to colours. Defaults to None.
"""
image_data = {
"image": image,
"title": title,
"color_map": color_map,
}
image_data = {"image": image, "title": title, "color_map": color_map}
self.images.append(image_data)

def generate(self) -> np.ndarray:
Expand All @@ -379,6 +336,7 @@ def generate(self) -> np.ndarray:
matplotlib.use("Agg")

self.figure, self.axis = plt.subplots(1, num_cols, figsize=figure_size)
self.figure.subplots_adjust(right=0.9)

axes = self.axis if isinstance(self.axis, np.ndarray) else np.array([self.axis])
for axis, image_dict in zip(axes, self.images, strict=True):
Expand Down

0 comments on commit 7ced8dd

Please sign in to comment.