Update anomalib/data/utils/boxes.py

Co-authored-by: Samet Akcay <[email protected]>
openvinotoolkit · Nov 28, 2022 · cec6138 · cec6138
1 parent 4c3a6b1
commit cec6138
Showing 1 changed file with 2 additions and 4 deletions.
diff --git 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: