From 97e6f4428f804e42d1d2cf413e20d25596c8d868 Mon Sep 17 00:00:00 2001 From: Vladimir Iglovikov Date: Sat, 1 Feb 2025 16:59:16 -0800 Subject: [PATCH] Cleanup --- albumentations/augmentations/transforms.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/albumentations/augmentations/transforms.py b/albumentations/augmentations/transforms.py index 7cc994f0a..f13d10306 100644 --- a/albumentations/augmentations/transforms.py +++ b/albumentations/augmentations/transforms.py @@ -5777,16 +5777,17 @@ def get_params_dependent_on_data( total_amount = self.py_random.uniform(*self.amount) salt_ratio = self.py_random.uniform(*self.salt_vs_pepper) - num_pixels = int(height * width * total_amount) + area = height * width + + num_pixels = int(area * total_amount) num_salt = int(num_pixels * salt_ratio) # Generate all positions at once - all_positions = np.arange(height * width) - noise_positions = self.random_generator.choice(all_positions, size=num_pixels, replace=False) + noise_positions = self.random_generator.choice(area, size=num_pixels, replace=False) # Create masks - salt_mask = np.zeros(height * width, dtype=bool) - pepper_mask = np.zeros(height * width, dtype=bool) + salt_mask = np.zeros(area, dtype=bool) + pepper_mask = np.zeros(area, dtype=bool) # Set salt and pepper positions salt_mask[noise_positions[:num_salt]] = True