gae/utils.py

import cv2
import numpy as np
from scipy import ndimage
import matplotlib.pyplot as plt

from skimage import exposure

import tensorflow as tf
import tensorflow_addons as tfa


IMAGE_SIZE = 224


def acc(y_true, y_pred, threshold=0.5):
    threshold = tf.cast(threshold, y_pred.dtype)
    y_pred = tf.cast(y_pred > threshold, y_pred.dtype)
    return tf.reduce_mean(tf.cast(tf.equal(y_true, y_pred), tf.float32))

def mae(y_true, y_pred):
    return tf.reduce_mean(tf.abs(y_true-y_pred))

def inv_ssim(y_true, y_pred):
    return 1 - tf.reduce_mean(tf.image.ssim(y_true, y_pred, 1.0))

def inv_msssim(y_true, y_pred):
    return 1 - tf.reduce_mean(tf.image.ssim_multiscale(y_true, y_pred, 1.0, filter_size=4))

def inv_msssim_l1(y_true, y_pred, alpha=0.8):
    return alpha*inv_msssim(y_true, y_pred) + (1-alpha)*mae(y_true, y_pred)

def inv_msssim_gaussian_l1(y_true, y_pred, alpha=0.8):
    l1_diff = tf.abs(y_true-y_pred)
    gaussian_l1 = tfa.image.gaussian_filter2d(l1_diff, filter_shape=(11, 11), sigma=1.5)
    return alpha*inv_msssim(y_true, y_pred) + (1-alpha)*gaussian_l1

def psnr(y_true, y_pred):
    return tf.reduce_mean(tf.image.psnr(y_true, y_pred, 1.0))


def show_image(image, title='Image', cmap_type='gray'):
    plt.imshow(image, cmap=cmap_type)
    plt.title(title)
    plt.axis('off')
    plt.show()


# đảo màu những ảnh bị ngược màu
def remove_negative(img):
  outside = np.mean(img[ : , 0])
  inside = np.mean(img[ : , int(IMAGE_SIZE / 2)])
  if outside < inside:
    return img
  else:
    return 1 - img

# lựa chọn tiền xử lý: ảnh gốc, Equalization histogram, CLAHE
def preprocess(img):
    img = remove_negative(img)

    img = exposure.equalize_hist(img)
    img = exposure.equalize_adapthist(img)
    img = exposure.equalize_hist(img)
    return img


# dilate contour
def dilate(mask_img):
    kernel_size = 2 * 22 + 1
    kernel = np.ones((kernel_size, kernel_size), dtype=np.uint8)
    return ndimage.binary_dilation(mask_img == 0, structure=kernel)


def normalize_tuple(value, n, name, allow_zero=False):
    """Transforms non-negative/positive integer/integers into an integer tuple.
    Args:
      value: The value to validate and convert. Could an int, or any iterable of
        ints.
      n: The size of the tuple to be returned.
      name: The name of the argument being validated, e.g. "strides" or
        "kernel_size". This is only used to format error messages.
      allow_zero: Default to False. A ValueError will raised if zero is received
        and this param is False.
    Returns:
      A tuple of n integers.
    Raises:
      ValueError: If something else than an int/long or iterable thereof or a
      negative value is
        passed.
    """
    error_msg = (
        f"The `{name}` argument must be a tuple of {n} "
        f"integers. Received: {value}"
    )

    if isinstance(value, int):
        value_tuple = (value,) * n
    else:
        try:
            value_tuple = tuple(value)
        except TypeError:
            raise ValueError(error_msg)
        if len(value_tuple) != n:
            raise ValueError(error_msg)
        for single_value in value_tuple:
            try:
                int(single_value)
            except (ValueError, TypeError):
                error_msg += (
                    f"including element {single_value} of "
                    f"type {type(single_value)}"
                )
                raise ValueError(error_msg)

    if allow_zero:
        unqualified_values = {v for v in value_tuple if v < 0}
        req_msg = ">= 0"
    else:
        unqualified_values = {v for v in value_tuple if v <= 0}
        req_msg = "> 0"

    if unqualified_values:
        error_msg += (
            f" including {unqualified_values}"
            f" that does not satisfy the requirement `{req_msg}`."
        )
        raise ValueError(error_msg)

    return value_tuple


def draw_points(image, points, color=None, random_color=False, same=True, thickness=1):
    if color is None and not random_color:
        color = (0, 255, 0)  # Màu mặc định là xanh lá cây (BGR)
    if random_color:
        color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))

    image = to_color(image)

    for point in points:
        if random_color and not same:
            color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))

        x, y = point
        image = cv2.circle(image, (x, y), thickness, color, -1)  # Vẽ điểm lên ảnh
    return image


def draw_lines(image, pairs, color=None, random_color=False, same=True, thickness=1):
    image_with_line = to_color(np.copy(image))

    if color is None and not random_color:
        color = (0, 255, 0)  # Màu mặc định là xanh lá cây (BGR)
    if random_color:
        color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))

    # Vẽ đường thẳng dựa trên danh sách các cặp điểm
    for pair in pairs:

        if random_color and not same:
            color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))

        start_point = pair[0]
        end_point = pair[1]
        image_with_line = cv2.line(image_with_line, start_point, end_point, color, thickness)
        image_with_line = cv2.circle(image_with_line, start_point, thickness + 1, color, -1)
        image_with_line = cv2.circle(image_with_line, end_point, thickness + 1, color, -1)

    return image_with_line

def to_color(image):
    if len(image.shape) == 3 and image.shape[-1] == 3:
        return image
    return cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)


def to_gray(image):
    if len(image.shape) == 3 and image.shape[-1] == 3:
        return cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    return image