#!/usr/bin/python3
import os
from typing import List, Tuple
import cv2
import numpy as np
from tensorflow.python.keras.models import Model as keras_model
[docs]class ImageUtils:
"""Util Class to hold all image processing methods.
Static methods are used to do some I/O tasks and processing images,
without instantiating the util class.
"""
[docs] @staticmethod
def mkdir_s(path: str) -> None:
"""Create directory in specified path, if not exists.
Parameters
----------
path: str
directory name to create
Returns
-------
None
Example
-------
unagi.utils.img_processing_utils.ImageUtils.mkdir_s(dir_name)
"""
if not os.path.exists(path):
os.makedirs(path)
[docs] @staticmethod
def normalize_in(img: np.ndarray) -> np.ndarray:
"""Normalize the input image to have all pixels in range 0 to 1.
Parameters
----------
img: numpy.ndarray
image array to normalize
Returns
-------
numpy.ndarray
normalized image array
Example
-------
unagi.utils.img_processing_utils.ImageUtils.normalize_in(img_array)
"""
img = img.astype(np.float32)
img /= 256.0
img -= 0.5
return img
[docs] @staticmethod
def normalize_gt(img: np.ndarray) -> np.ndarray:
"""Normalize the gt image to have all pixels in range 0 to 1.
Parameters
----------
img: numpy.ndarray
image array to normalize
Returns
-------
numpy.ndarray
normalized image array
Example
-------
unagi.utils.img_processing_utils.ImageUtils.normalize_gt(img_array)
"""
img = img.astype(np.float32)
img /= 255.0
return img
[docs] @staticmethod
def add_border(
img: np.ndarray, size_x: int = 128, size_y: int = 128
) -> Tuple[np.ndarray, int, int]:
"""Add border to image based on the inputs
Parameters
----------
img: numpy.ndarray
image array to add border
size_x: int
width for image part (deafult is `128`).
size_y: int
height for image part (deafult is `128`).
Returns
-------
Tuple[np.ndarray, int, int]
image array with border
border value added on width
border value added on height
Example
-------
unagi.utils.img_processing_utils.ImageUtils.add_border(img_array, 128, 128)
"""
max_y, max_x = img.shape[:2]
border_y = 0
if max_y % size_y != 0:
border_y = (size_y - (max_y % size_y) + 1) // 2
img = cv2.copyMakeBorder(
img,
border_y,
border_y,
0,
0,
cv2.BORDER_CONSTANT,
value=[255, 255, 255],
)
border_x = 0
if max_x % size_x != 0:
border_x = (size_x - (max_x % size_x) + 1) // 2
img = cv2.copyMakeBorder(
img,
0,
0,
border_x,
border_x,
cv2.BORDER_CONSTANT,
value=[255, 255, 255],
)
return img, border_y, border_x
[docs] @staticmethod
def split_img(
img: np.ndarray, size_x: int = 128, size_y: int = 128
) -> List[np.ndarray]:
"""Split image to parts (little images).
Parameters
----------
img: numpy.ndarray
image array to split
size_x: int
width for image part (deafult is `128`).
size_y: int
height for image part (deafult is `128`).
Returns
-------
List[numpy.ndarray]
list of numpy.ndarray of image parts
Note
----
Walk through the whole image by the window of size size_x * size_y
without overlays and save all parts in list.
Images sizes should divide window sizes.
Example
-------
unagi.utils.img_processing_utils.ImageUtils.split_img(img_array, 128, 128)
"""
max_y, max_x = img.shape[:2]
parts = []
curr_y = 0
# TODO: rewrite with generators.
while (curr_y + size_y) <= max_y:
curr_x = 0
# fmt: off
while (curr_x + size_x) <= max_x:
parts.append(img[curr_y : curr_y + size_y, curr_x : curr_x + size_x]) # fmt: skip # noqa
curr_x += size_x
curr_y += size_y
# fmt: on
return parts
[docs] @staticmethod
def combine_imgs(imgs: List[np.ndarray], max_y: int, max_x: int) -> np.ndarray:
"""Combine image parts to one big image.
Parameters
----------
img: List[numpy.ndarray]
list image arraies to combine
max_y: int
width for image part (deafult is `128`).
max_x: int
height for image part (deafult is `128`).
Returns
-------
numpy.ndarray
numpy.ndarray of combined image
Note
----
Walk through list of images and create from them one big image
with sizes max_x * max_y.
If border_x and border_y are non-zero,
they will be removed from created image.
The list of images should contain data in the following order:
from left to right, from top to bottom.
Example
-------
unagi.utils.img_processing_utils.ImageUtils.combine_imgs(
img_array_list, 128, 128)
"""
img = np.zeros((max_y, max_x), np.float)
size_y, size_x = imgs[0].shape
curr_y = 0
i = 0
# TODO: rewrite with generators.
while (curr_y + size_y) <= max_y:
curr_x = 0
while (curr_x + size_x) <= max_x:
# fmt: off
try:
img[curr_y : curr_y + size_y, curr_x : curr_x + size_x] = imgs[i] # fmt: skip # noqa
except Exception:
i -= 1
i += 1
curr_x += size_x
# fmt: on
curr_y += size_y
return img
[docs] @staticmethod
def shuffle_imgs(dname: str) -> None:
"""Shuffle input and ground-truth images.
(actual, if You are using different datasets as one).
Parameters
----------
dname: str
directory name with image to shuffle
Returns
-------
None
Example
-------
unagi.utils.img_processing_utils.ImageUtils.shuffle_imgs(images_dir)
"""
dir_in = os.path.join(dname, "in")
dir_gt = os.path.join(dname, "gt")
n = len(os.listdir(dir_in))
np.random.seed()
for i in range(n):
j = i
while j == i:
j = np.random.randint(0, n)
os.rename(
os.path.join(dir_in, str(i) + "_in.png"),
os.path.join(dir_in, "tmp_in.png"),
)
os.rename(
os.path.join(dir_in, str(j) + "_in.png"),
os.path.join(dir_in, str(i) + "_in.png"),
)
os.rename(
os.path.join(dir_in, "tmp_in.png"),
os.path.join(dir_in, str(j) + "_in.png"),
)
os.rename(
os.path.join(dir_gt, str(i) + "_gt.png"),
os.path.join(dir_gt, "tmp_gt.png"),
)
os.rename(
os.path.join(dir_gt, str(j) + "_gt.png"),
os.path.join(dir_gt, str(i) + "_gt.png"),
)
os.rename(
os.path.join(dir_gt, "tmp_gt.png"),
os.path.join(dir_gt, str(j) + "_gt.png"),
)
[docs] @staticmethod
def process_with_unagi(
img: np.ndarray, model: keras_model, batchsize: int
) -> np.ndarray:
"""Split image to 128x128 parts and run U-net for every part.
Parameters
----------
img: numpy.ndarray
image array to preprocess
model: keras_model
keras model
batchsize: int
batchsize to use with the model
Returns
-------
numpy.ndarray
image array after preprocessing
See Also
--------
ImageUtils.add_border(), ImageUtils.normalize_in(),
ImageUtils.split_img(), ImageUtils.combine_imgs()
Example
-------
unagi.utils.img_processing_utils.ImageUtils.process_with_unagi(
process_unet_img, model)
"""
img, border_y, border_x = ImageUtils.add_border(img)
img = ImageUtils.normalize_in(img)
parts = ImageUtils.split_img(img)
parts_array = np.array(parts)
parts_array.shape = (
parts_array.shape[0],
parts_array.shape[1],
parts_array.shape[2],
1,
)
parts = model.predict(parts_array, batchsize, verbose=1)
tmp = []
for part in parts:
part.shape = (128, 128)
tmp.append(part)
parts = tmp
img = ImageUtils.combine_imgs(parts, img.shape[0], img.shape[1])
# fmt: off
img = img[
border_y : img.shape[0] - border_y, border_x : img.shape[1] - border_x # fmt: skip # noqa
]
# fmt: on
img = img * 255.0
img = img.astype(np.uint8)
return img
[docs] @staticmethod
def postprocess_img(img: np.ndarray) -> np.ndarray:
"""Apply Otsu threshold to image.
Parameters
----------
img: numpy.ndarray
image array to postprocess
Returns
-------
numpy.ndarray
postprocessed image array
Example
-------
unagi.utils.img_processing_utils.ImageUtils.postprocess_img(img_array)
"""
_, img = cv2.threshold(img, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
return img
[docs] @staticmethod
def binarize_img(
img: np.ndarray, model: keras_model, batchsize: int = 2
) -> np.ndarray:
"""Binarize image, using U-net, Otsu, bottom-hat transform etc.
Parameters
----------
img: numpy.ndarray
image array to preprocess
model: keras_model
keras model
batchsize: int, optional
batchsize to use with the model (default is `2`)
Returns
-------
numpy.ndarray
image array after binarizing
See Also
--------
ImageUtils.process_with_unagi(), ImageUtils.postprocess_img()
Example
-------
unagi.utils.img_processing_utils.ImageUtils.binarize_img(
process_unet_img, model)
"""
img = ImageUtils.process_with_unagi(img, model, batchsize)
img = ImageUtils.postprocess_img(img)
return img