Some noises

some_noises.py

def salt_and_pepper(img, noise_probability=1, noise_density=0.2, salt_ratio=0.1):
    """
    Given a numpy image, return a salt and pepper noised image.

    :param img: Image in which noise will be added. The first dimension will be taken as the batch element index.
    :param noise_probability: Probability of adding the noise (by default 1).
    :param noise_density: Amount of noise in the image (as a percentage).
    :param salt_ratio: Ratio of salt (vs pepper) added in the image.
    :return: Noised batch of images.
    """
    batch_size = img.shape[0]
    output = np.copy(img).astype(np.uint8)
    noise_density = np.random.uniform(0, noise_density)
    for i in range(batch_size):
        r = random.uniform(0, 1)  # Random number
        if noise_probability >= r:  # Inside the probability
            black_dots = (np.random.uniform(0, 1, output[i, :, :].shape) > noise_density * (1 - salt_ratio)).astype(
                np.uint8)
            white_dots = 1 - (np.random.uniform(0, 1, output[i, :, :].shape) > noise_density * salt_ratio).astype(
                np.uint8)
            output[i, :, :] = np.logical_and(output[i, :, :], black_dots)
            output[i, :, :] = np.logical_or(output[i, :, :], white_dots)
    return output
  
  
  
#   El ruido de Agos (la funcion de abajo es un wrapper para el batch)

def draw_lines_single(x, swap_prob=0.6):
    noised = x  # Output starts from input
    width, height = x.shape

    n_10 = int(width * 10 / 1024)
    n_30 = int(width * 30 / 1024)
    n_60 = int(width * 60 / 1024)
    n_66 = int(width * 66 / 1024)
    n_80 = int(width * 80 / 1024)
    n_100 = int(width * 100 / 1024)
    n_150 = int(width * 150 / 1024)
    n_900 = int(width * 900 / 1024)
    n_950 = int(width * 950 / 1024)
    n_1024 = int(width)

    num_lines = np.random.randint(1, n_100)  # Pick a number N between 1-100
    for line in range(num_lines):  # Put N random white circles in output
        pos_x, pos_y = np.random.randint(n_100, n_950), np.random.randint(n_100, n_950)
        radius = np.random.randint(n_60)
        cv.circle(noised, (pos_x, pos_y), radius, color=(1, 1, 1), thickness=-1)  # Draw circle in noised

    k = np.random.randint(n_10, n_30)  # Square kernel size
    kernel = np.ones((k, k), np.uint8)  # Create square kernel
    noised = cv.morphologyEx(noised, cv.MORPH_CLOSE, kernel)  # Morphological closing with random sized square kernel

    num_lines = np.random.randint(1, n_80)
    for line in range(num_lines):  # Put N random black circles in output
        pos_x, pos_y = np.random.randint(n_150, n_900), np.random.randint(n_150, n_900)
        radius = np.random.randint(n_30)
        cv.circle(noised, (pos_x, pos_y), radius, color=(0, 0, 0), thickness=-1)  # Draw circle in noised

    if np.random.uniform(0, 1) > 0.3:  # Check if inside probability
        swap_map = np.random.choice(range(2), size=(n_1024, n_66),
                                    p=[1 - swap_prob, swap_prob])  # Slice of the image of rand. pixels
        h_idx, w_idx = np.where(swap_map == 1)  # Check which of those pixels is one
        w_idx = 15 * (w_idx + 1)
        bor = 15
        for b in range(bor):
            noised[h_idx, w_idx + b - bor] = noised[h_idx, w_idx - (1 + b)]
            noised[h_idx, w_idx - (1 + b)] = noised[h_idx, w_idx + b - bor]

    if np.random.uniform(0, 1) > 0.4:
        kernel = np.ones((1, 7), np.uint8)
        noised = cv.morphologyEx(noised, cv.MORPH_CLOSE, kernel)  # Closing with horizontal kernel

    return noised


def draw_lines(x, swap_prob=0.6):
    if len(x.shape) == 2:
        return draw_lines_single(x, swap_prob)

    one_channel = False
    if len(x.shape) == 3:
        one_channel = True
        x = np.expand_dims(x, 1)
    batch_size, channels, height, width = x.shape
    noised = x

    for i in range(batch_size):
        for j in range(channels):
            noised[i, j, :, :] = draw_lines_single(x[i, j, :, :], swap_prob)
    if one_channel:
        return noised[:, 0, :, :]
    else:
        return noised
      
      
# En 3D

def shape_3d(center, size, image_size, shape="sphere"):
    """ Generate a 3D numpy sphere or cube given the center, size and image size. 
    
    It creates a mask of the shape usingthe p-norm concept.

    :param center: array with the coordinates center of the shape.
    :param size: single number that represents the size of the shape in each dimension.
    :param image_size: size of the output array.
    :param shape: shape to return.
    :return:
    """
    if type(image_size) == sitk.Image:
        image_size = sitk.GetArrayFromImage(image_size).shape

    if shape in ["circle", "sphere"]:
        ord = 2
    elif shape in ["square", "box", "cube"]:
        ord = np.inf
    else:
        print(
            "Shape {} is not supported. Setting shape as sphere".format(shape)
        )
        ord = 2
    distance = np.linalg.norm(
        np.subtract(np.indices(image_size).T, np.asarray(center)),
        axis=len(center),
        ord=ord,
    )
    shape_np = 1 - np.ones(image_size).T * (distance <= size)
    return shape_np.T