import cv2
import numpy as np


def edge_prewitt(image):
    """
    :param image:
    :return:
        Prewitt:
              [ 1,  0, -1]        [ 1,  1,  1]
        G_x = [ 1,  0, -1]  G_y = [ 1,  0, -1]
              [ 1,  0, -1]        [-1,  0, -1]
        Prewitt = threshold(max(G_x * I, G_y * I))
    """
    g_x = np.array([[1, 0, -1], [1, 0, -1], [1, 0, -1]], dtype=int)
    g_y = np.array([[1, 1, 1], [0, 0, 0], [-1, -1, -1]], dtype=int)
    x = cv2.filter2D(image, cv2.CV_16S, g_x)
    y = cv2.filter2D(image, cv2.CV_16S, g_y)
    absX = cv2.convertScaleAbs(x)
    absY = cv2.convertScaleAbs(y)
    result = cv2.max(absX, absY)
    _, result = cv2.threshold(result, 0, 255, cv2.THRESH_OTSU)
    return result


def edge_sobel(image, kernel_size=3):
    x = cv2.Sobel(image, cv2.CV_16S, 1, 0, ksize=kernel_size)
    y = cv2.Sobel(image, cv2.CV_16S, 0, 1, ksize=kernel_size)
    absX = cv2.convertScaleAbs(x)
    absY = cv2.convertScaleAbs(y)
    result = cv2.addWeighted(absX, 0.5, absY, 0.5, 0)
    _, result = cv2.threshold(result, 0, 255, cv2.THRESH_OTSU)
    return result


# 1. 读取图像文件
image_path = './images/lena.jpg'
image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)

# 2. Prewitt
prewitt = edge_prewitt(image)

# 3. Sobel
sobel = edge_sobel(image, kernel_size=3)

# 4. Canny
canny = cv2.Canny(image, 50, 150)

# 5. 显示
cv2.imshow('Image', image)
cv2.imshow('Prewitt', prewitt)
cv2.imshow('Sobel', sobel)
cv2.imshow('Canny', canny)
cv2.waitKey(0)