diff --git a/monai/transforms/utils.py b/monai/transforms/utils.py index ffbd91690c..23c6bd100a 100644 --- a/monai/transforms/utils.py +++ b/monai/transforms/utils.py @@ -344,15 +344,16 @@ def _correct_centers( return center_ori centers = [] + fg_indices, bg_indices = np.asarray(fg_indices), np.asarray(bg_indices) + if fg_indices.size == 0 and bg_indices.size == 0: + raise ValueError("No sampling location available.") - if not len(fg_indices) or not len(bg_indices): - if not len(fg_indices) and not len(bg_indices): - raise ValueError("No sampling location available.") + if fg_indices.size == 0 or bg_indices.size == 0: warnings.warn( f"N foreground {len(fg_indices)}, N background {len(bg_indices)}," "unable to generate class balanced samples." ) - pos_ratio = 0 if not len(fg_indices) else 1 + pos_ratio = 0 if fg_indices.size == 0 else 1 for _ in range(num_samples): indices_to_use = fg_indices if rand_state.rand() < pos_ratio else bg_indices