Skip to content

Commit

Permalink
Cleanup
Browse files Browse the repository at this point in the history
  • Loading branch information
ternaus committed Feb 2, 2025
1 parent 3bd82e3 commit 97e6f44
Showing 1 changed file with 6 additions and 5 deletions.
11 changes: 6 additions & 5 deletions albumentations/augmentations/transforms.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 97e6f44

Please sign in to comment.