From cec6138f09a3a96f3c3edb6d92f2105f05e9126b Mon Sep 17 00:00:00 2001 From: Dick Ameln Date: Mon, 28 Nov 2022 13:23:55 +0100 Subject: [PATCH] Update anomalib/data/utils/boxes.py Co-authored-by: Samet Akcay --- anomalib/data/utils/boxes.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/anomalib/data/utils/boxes.py b/anomalib/data/utils/boxes.py index df9a454e20..5601478c23 100644 --- a/anomalib/data/utils/boxes.py +++ b/anomalib/data/utils/boxes.py @@ -24,10 +24,8 @@ def masks_to_boxes(masks: Tensor) -> List[Tensor]: masks = masks.view((-1, 1) + masks.shape[-2:]) # reshape to (B, 1, H, W) masks = masks.float() - if masks.is_cuda: - batch_comps = connected_components_gpu(masks).squeeze(1) - else: - batch_comps = connected_components_cpu(masks).squeeze(1) + connected_components = connected_components_gpu if masks.is_cuda else connected_components_cpu + batch_comps = connected_components(masks).squeeze(1) batch_boxes = [] for im_comps in batch_comps: