visiontoolset.py

import math
from pathlib import Path

import cv2 as cv2
import numpy as np
from PIL import ImageFont, ImageDraw, Image
from skimage import morphology
from skimage.util import img_as_ubyte


def get_bbox(image):
    """ Get the bounding box of the full character set in the image and return the coordinates

    :param image: image with characters
    :return: coordinates of the bounding box, min_x, min_y, max_x, max_y
    """
    assert image is not None, "input image is None"

    min_x, min_y = np.inf, np.inf
    max_x, max_y = -np.inf, -np.inf

    # Apply binary threshold and find contours
    _, thresh = cv2.threshold(image, 127, 255, cv2.THRESH_BINARY)
    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # Loop through all contours to find the bounding boxes
    for cnt in contours:
        x, y, w, h = cv2.boundingRect(cnt)
        min_x, min_y = min(min_x, x), min(min_y, y)
        max_x, max_y = max(max_x, x + w), max(max_y, y + h)
    return min_x, min_y, max_x, max_y


def skeletonize(image):
    """ Perform a skeletonization on the image

    :param image: input image
    :return: skeleton image
    """
    assert image is not None, "input image is None"

    thin_image = morphology.thin(image)
    thin_image = img_as_ubyte(thin_image)
    return thin_image


def get_points(image):
    """ make a list of all the points in the skeleton image

    :param image: image
    :return: list of points tuple (x,y)
    """
    assert image is not None, "Skeleton image is None"

    points = []
    for r in range(0, image.shape[0]):
        for c in range(0, image.shape[1]):
            if image[r, c] == 255:
                points.append((c, r))
    return points


def character_to_dots(character='a', font_path=r"C:\Windows\Fonts\Arial.ttf", font_size=64):
    """ Create an image from a given font and character

    :param character: character to create
    :param font_path: path to the font file including font name like r"C:\Windows\Fonts\Arial.ttf"
    :param font_size: font size
    :return: images with the character and its skeleton
    """
    assert Path(font_path).exists(), "Font file does not exist"
    assert character is not None, "Character is None"

    nr_of_char = len(character)
    font = ImageFont.truetype(font_path, font_size)
    img_font = Image.new('L', (int(nr_of_char * font_size * 1.5), int(font_size * 1.5)), color=0)
    d = ImageDraw.Draw(img_font)
    d.text((0, 0), character, font=font, fill=255)
    img_font = np.array(img_font)  # convert to numpy array for opencv

    # Get image bounding box and crop to it but only if it is not empty (space character)
    min_x, min_y, max_x, max_y = get_bbox(img_font)
    if min_x != np.inf:
        img_text = img_font[min_y:max_y, min_x:max_x]
    else:
        img_text = img_font

    # create a skeleton of the image
    img_skel = skeletonize(img_text)

    return img_text, img_skel


def rotate_image(image, angle):
    """ Rotate the image by the given angle in degrees

    :param image: input image
    :param angle: angle in degrees
    :return: rotated image
    """

    assert image is not None, "input image is None"

    height, width = image.shape[:2]
    angle_rad = math.radians(angle)
    center_x, center_y = width // 2, height // 2

    # Calculate the new image width and height
    new_width = abs(width * math.cos(angle_rad)) + abs(height * math.sin(angle_rad))
    new_height = abs(height * math.cos(angle_rad)) + abs(width * math.sin(angle_rad))
    new_width, new_height = int(new_width), int(new_height)

    # Calculate the rotation matrix, considering the new image size
    rotation_matrix = cv2.getRotationMatrix2D((center_x, center_y), angle, 1.0)

    # Adjust the rotation matrix to account for translation
    rotation_matrix[0, 2] += (new_width / 2) - center_x
    rotation_matrix[1, 2] += (new_height / 2) - center_y

    # Perform the rotation
    rotated_img = cv2.warpAffine(image, rotation_matrix, (new_width, new_height))

    return rotated_img


def resize_and_centre_image(image, sample_size=(64, 64)):
    """ Creates a new image with a new size and centers the input image onto that image

    :param image: input image
    :param size: new size
    :return: resized image
    """
    assert image is not None, "input image is None"
    assert sample_size is not None, "size is None"

    diff = np.subtract(sample_size, image.shape[:2]) // 2
    start_y = max(diff[0], 0)
    start_x = max(diff[1], 0)
    stop_y = min(diff[0] + image.shape[0], sample_size[0])
    stop_x = min(diff[1] + image.shape[1], sample_size[1])

    if len(image.shape) == 2:  # for gray images
        diff = np.subtract(sample_size, image.shape) // 2
        new_img = np.zeros(sample_size, dtype=np.uint8)
        new_img[start_y:stop_y, start_x:stop_x] = image[max(-diff[0], 0):max(-diff[0], 0) + stop_y - start_y,
                                                        max(-diff[1], 0):max(-diff[1], 0) + stop_x - start_x]

    else:  # for color images
        new_img = np.zeros(sample_size + (3,), dtype=np.uint8)
        new_img[start_y:stop_y, start_x:stop_x, :] = image[max(-diff[0], 0):max(-diff[0], 0) + stop_y - start_y,
                                                           max(-diff[1], 0):max(-diff[1], 0) + stop_x - start_x, :]
    return new_img