Merge pull request #3 from ErickOF/feature-datagen_salt_and_pepper_noise

The datagen tool add salt and pepper noise to grayscale image

Author: ErickOF

.gitignore

@@ -221,4 +221,5 @@ pyrightconfig.json
 
 # Ignore generate images from datagen tool
 tools/datagen/src/imgs/*_grayscale_image*
+tools/datagen/src/imgs/*_noisy_image*
 tools/datagen/src/imgs/*_sobel_*

tools/datagen/README.md

@@ -22,19 +22,25 @@ There are two versions of the script, a JupyterNotebook and a normal Python
 script. For both of them, in the main function, the variables `img_name` and
 `img_ext` must be updated. The `img_name` is only the name of the image that
 is saved in `tools/datagen/src/imgs` and the `img_ext` is the extension of the
-image. For instance:
+image. The `salt_probability` is used to set the probability of getting a salt
+pixel in the image, while `pepper_probability` is used for the pepper pixels.
+These two variables can be set to 0.0 to generate an image without noise. For
+instance:
 
 ```python
+salt_probability: float = 0.1
+salt_probability: float = 0.1
 img_name: str = 'car'
 img_ext: str = 'jpg
 ```
 
 will load `tools/datagen/src/imgs/car.jpg` image and generate:
 
 * tools/datagen/src/imgs/car\_grayscale\_image.jpg
-* tools/datagen/src/imgs/car\_sobel\_combined\_result.jpg
+* tools/datagen/src/imgs/car\_noisy\_image.jpg
 * tools/datagen/src/imgs/car\_sobel\_x\_result.jpg
 * tools/datagen/src/imgs/car\_sobel\_y\_result.jpg
+* tools/datagen/src/imgs/car\_sobel\_combined\_result.jpg
 
 These generated images are ignored and must not be committed to the repository.
 

tools/datagen/datagen.ipynb

@@ -2,32 +2,56 @@
 """
 Script to generate the data required for SystemC testbench.
 """
-from typing import Tuple
+from typing import List, Tuple
 
 import cv2
 import numpy as np
 import matplotlib.pyplot as plt
 
 
-def apply_sobel_filter(image_path: str
+def apply_sobel_filter(
+        image_path: str,
+        salt_prob: float = 0.0,
+        pepper_prob: float = 0.0,
     )-> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
     """Applies the Sobel filter in X and Y to a input image and returns each
     step.
 
     Args:
         image_path (str): Path to the input image
+        salt_prob (float): Probability of changing a pixel to salt. Value from
+            0.0 to 1.0. Defaults to 0.0.
+        pepper_prob (float): Probability of changing a pixel to pepper. Value
+            from 0.0 to 1.0. Defaults to 0.0.
 
     Returns:
-        Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: A tuple
-            containing the gray image, the sobel in x, the sobel in y, and the
-            sobel combining x and y.
+        Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: A
+            tuple containing the grayscale image, the grayscale image with salt
+            and pepper noise, the sobel in x, the sobel in y, and the sobel
+            combining x and y.
     """
     # Read the image
     image: np.ndarray = cv2.imread(image_path)
 
     # BGR to grayscale
     gray_img: np.ndarray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
 
+    # Noisy image
+    noisy_img: np.ndarray = np.copy(gray_img)
+
+    salt_num: int = int(np.ceil(salt_prob * noisy_img.size))
+    pepper_num: int = int(np.ceil(pepper_prob * noisy_img.size))
+
+    salt_coords: List[np.ndarray] = [
+        np.random.randint(0, i, salt_num) for i in image.shape
+    ]
+    pepper_coords: List[np.ndarray] = [
+        np.random.randint(0, i, pepper_num) for i in image.shape
+    ]
+
+    noisy_img[salt_coords[0], salt_coords[1]] = 255
+    noisy_img[pepper_coords[0], pepper_coords[1]] = 0
+
     # Apply Sobel in x direction
     sobel_x_img: np.ndarray = cv2.Sobel(gray_img, cv2.CV_64F, 1, 0, ksize=3)
     sobel_x_uint8: np.ndarray = np.uint8(np.absolute(sobel_x_img))
@@ -40,49 +64,69 @@ def apply_sobel_filter(image_path: str
     sobel_combined_img: np.ndarray = cv2.magnitude(sobel_x_img, sobel_y_img)
 
     # Convert to 8-bit image
-    sobel_combined_uint8: np.ndarray = np.uint8(np.absolute(sobel_combined_img))
+    sobel_combined_uint8: np.ndarray = np.uint8(
+        np.absolute(sobel_combined_img)
+    )
 
     # Display the results
     plt.figure(figsize=(10, 8))
 
-    plt.subplot(2, 2, 1)
+    plt.subplot(2, 3, 1)
     plt.title('Original Image')
     plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
     plt.axis('off')
 
-    plt.subplot(2, 2, 2)
+    plt.subplot(2, 3, 2)
     plt.title('Grayscale Image')
     plt.imshow(gray_img, cmap='gray')
     plt.axis('off')
 
-    plt.subplot(2, 2, 3)
+    plt.subplot(2, 3, 3)
+    plt.title('Noisy Image')
+    plt.imshow(noisy_img, cmap='gray')
+    plt.axis('off')
+
+    plt.subplot(2, 3, 4)
     plt.title('Sobel X')
     plt.imshow(np.uint8(np.absolute(sobel_x_img)), cmap='gray')
     plt.axis('off')
 
-    plt.subplot(2, 2, 4)
+    plt.subplot(2, 3, 5)
     plt.title('Sobel Y')
     plt.imshow(np.uint8(np.absolute(sobel_y_img)), cmap='gray')
     plt.axis('off')
 
-    plt.figure(figsize=(5, 5))
+    plt.subplot(2, 3, 6)
     plt.title('Sobel Combined')
     plt.imshow(sobel_combined_uint8, cmap='gray')
     plt.axis('off')
 
     plt.show()
 
-    return (gray_img, sobel_x_uint8, sobel_y_uint8, sobel_combined_uint8)
+    return (gray_img, noisy_img, sobel_x_uint8, sobel_y_uint8,
+            sobel_combined_uint8)
 
 
 if __name__ == '__main__':
+    salt_probability: float = 0.1
+    pepper_probability: float = 0.1
     img_name: str = 'car'
     img_ext: str = 'jpg'
     img_path: str = f'src/imgs/{img_name}.{img_ext}'
-    gray, sobel_x, sobel_y, sobel_combined = apply_sobel_filter(img_path)
+    gray, noisy, sobel_x, sobel_y, sobel_combined = \
+        apply_sobel_filter(img_path, salt_probability, pepper_probability)
 
     cv2.imwrite(f'src/imgs/{img_name}_grayscale_image.{img_ext}', gray)
-    cv2.imwrite(f'src/imgs/{img_name}_sobel_x_result.{img_ext}', np.uint8(np.absolute(sobel_x)))
-    cv2.imwrite(f'src/imgs/{img_name}_sobel_y_result.{img_ext}', np.uint8(np.absolute(sobel_y)))
-    cv2.imwrite(f'src/imgs/{img_name}_sobel_combined_result.{img_ext}', sobel_combined)
-
+    cv2.imwrite(f'src/imgs/{img_name}_noisy_image.{img_ext}', noisy)
+    cv2.imwrite(
+        f'src/imgs/{img_name}_sobel_x_result.{img_ext}',
+        np.uint8(np.absolute(sobel_x))
+    )
+    cv2.imwrite(
+        f'src/imgs/{img_name}_sobel_y_result.{img_ext}',
+        np.uint8(np.absolute(sobel_y))
+    )
+    cv2.imwrite(
+        f'src/imgs/{img_name}_sobel_combined_result.{img_ext}',
+        sobel_combined
+    )