# 对图片分别使用Roberts、Sobel和Prewitt三种边缘检测方式，完成边缘检测并使用matplotlib展示结果

import cv2 as cv
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import rcParams

# 导入图片
img = cv.imread('img.jpg', cv.IMREAD_GRAYSCALE)

# 设置中文字体
rcParams['font.sans-serif'] = ['SimHei']
rcParams['axes.unicode_minus'] = False

# Roberts 边缘检测
roberts_kernel_x = np.array([[1, 0], [0, -1]], dtype=np.float32)
roberts_kernel_y = np.array([[0, 1], [1, 0]], dtype=np.float32)
roberts_x = cv.filter2D(img.astype(np.float32), -1, roberts_kernel_x)
roberts_y = cv.filter2D(img.astype(np.float32), -1, roberts_kernel_y)
roberts = cv.magnitude(roberts_x, roberts_y)

# Sobel 边缘检测
sobel_X = cv.Sobel(img, cv.CV_64F, 1, 0, ksize=3)
Sobel_y = cv.Sobel(img, cv.CV_64F, 0, 1, ksize=3)
sobel = cv.magnitude(sobel_X, Sobel_y)

# Prewitt边缘检测
prewitt_kernel_x = np.array([[1, 0, -1], [1, 0, -1], [1, 0, -1]], dtype=np.float32)
prewitt_kernel_y = np.array([[1, 1, 1], [0, 0, 0], [-1, -1, -1]], dtype=np.float32)
prewitt_x = cv.filter2D(img.astype(np.float32), -1, prewitt_kernel_x)
prewitt_y = cv.filter2D(img.astype(np.float32), -1, prewitt_kernel_y)
prewitt = cv.magnitude(prewitt_x, prewitt_y)

# 使用 matplotlib 展示
plt.figure(figsize=(12, 8))
plt.subplot(2, 2, 1)
plt.title('原图')
plt.imshow(img, cmap='gray')
plt.axis('off')

plt.subplot(2, 2, 2)
plt.title('Roberts 边缘检测')
plt.imshow(roberts, cmap='gray')
plt.axis('off')

plt.subplot(2, 2, 3)
plt.title('Sobel 边缘检测')
plt.imshow(sobel, cmap='gray')
plt.axis('off')

plt.subplot(2, 2, 4)
plt.title('Prewitt 边缘检测')
plt.imshow(prewitt, cmap='gray')
plt.axis('off')

plt.tight_layout()
plt.show()