Spaces:
Runtime error
Runtime error
File size: 3,698 Bytes
c19ca42 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 |
from enum import Enum
from typing import Callable, List
import numpy as np
from ..utils.utils import get_h_w_c
from .image_utils import as_target_channels
def __add_noises(
image: np.ndarray,
noise_gen: Callable[[int, int], List[np.ndarray]],
combine: Callable[[np.ndarray, List[np.ndarray]], np.ndarray],
) -> np.ndarray:
img = image
h, w, c = get_h_w_c(img)
assert c != 2, "Noise cannot be added to 2-channel images."
if c > 3:
img = img[:, :, :3]
noises = noise_gen(h, w)
assert len(noises) > 0
max_channels = min(c, 3)
for n in noises:
noise_channels = get_h_w_c(n)[2]
assert noise_channels in (1, 3), "Noise must be a grayscale or RGB image."
max_channels = max(max_channels, noise_channels)
noises = [as_target_channels(n, max_channels) for n in noises]
img = as_target_channels(img, max_channels)
result = combine(img, noises)
if c > 3:
result = np.concatenate([result, image[:, :, 3:]], axis=2)
return np.clip(result, 0, 1)
def __add_noise(
image: np.ndarray,
noise_gen: Callable[[int, int], np.ndarray],
combine: Callable[[np.ndarray, np.ndarray], np.ndarray] = lambda i, n: i + n,
) -> np.ndarray:
return __add_noises(
image,
lambda h, w: [noise_gen(h, w)],
lambda i, n: combine(i, n[0]),
)
class NoiseColor(Enum):
RGB = "rgb"
GRAY = "gray"
@property
def channels(self):
return 3 if self is NoiseColor.RGB else 1
# Applies gaussian noise to an image
def gaussian_noise(
image: np.ndarray,
amount: float,
noise_color: NoiseColor,
seed: int = 0,
) -> np.ndarray:
rng = np.random.default_rng(seed)
return __add_noise(
image,
lambda h, w: rng.normal(0, amount, (h, w, noise_color.channels)).astype(
np.float32
),
)
# Applies uniform noise to an image
def uniform_noise(
image: np.ndarray,
amount: float,
noise_color: NoiseColor,
seed: int = 0,
) -> np.ndarray:
rng = np.random.default_rng(seed)
return __add_noise(
image,
lambda h, w: rng.uniform(-amount, amount, (h, w, noise_color.channels)).astype(
np.float32
),
)
# Applies salt and pepper noise to an image
def salt_and_pepper_noise(
image: np.ndarray,
amount: float,
noise_color: NoiseColor,
seed: int = 0,
) -> np.ndarray:
def gen_noise(h: int, w: int):
rng = np.random.default_rng(seed)
noise_c = noise_color.channels
amt = amount / 2
pepper = rng.choice([0, 1], (h, w, noise_c), p=[amt, 1 - amt]).astype(np.uint8)
salt = rng.choice([0, 1], (h, w, noise_c), p=[1 - amt, amt]).astype(np.uint8)
return [pepper, salt]
def combine(i: np.ndarray, n: List[np.ndarray]):
pepper, salt = n
return np.where(salt == 1, 1, np.where(pepper == 0, 0, i))
return __add_noises(image, gen_noise, combine)
# Applies poisson noise to an image
def poisson_noise(
image: np.ndarray,
amount: float,
noise_color: NoiseColor,
seed: int = 0,
) -> np.ndarray:
rng = np.random.default_rng(seed)
return __add_noise(
image,
lambda h, w: rng.poisson(amount, (h, w, noise_color.channels)).astype(np.uint8),
)
# Applies speckle noise to an image
def speckle_noise(
image: np.ndarray,
amount: float,
noise_color: NoiseColor,
seed: int = 0,
) -> np.ndarray:
rng = np.random.default_rng(seed)
return __add_noise(
image,
lambda h, w: rng.normal(0, amount, (h, w, noise_color.channels)).astype(
np.float32
),
lambda i, n: i + i * n,
)
|