import numpy as np
import cv2


def canny(image, min_threshold, max_threshold):
    # 1. 高斯滤波器，去噪声
    kernel_size = 5
    sigma = 1.4
    blur = cv2.GaussianBlur(image, (kernel_size, kernel_size), sigma)

    # 2. 计算梯度值和方向
    dx = cv2.Sobel(blur, cv2.CV_64F, 1, 0)
    dy = cv2.Sobel(blur, cv2.CV_64F, 0, 1)
    magnitude = np.sqrt(dx * dx + dy * dy)  # 计算梯度幅值
    angle = np.arctan2(dy, dx) * 180 / np.pi  # 计算梯度方向
    angle[angle < 0] += 180  # 将梯度方向转换为0-180度之间

    # 3. 非极大值抑制
    height, width = magnitude.shape
    new_magnitude = np.zeros((height, width))
    for i in range(1, height - 1):
        for j in range(1, width - 1):
            if (0 <= angle[i, j] < 22.5) or (157.5 <= angle[i, j] <= 180):
                if (magnitude[i, j] >= magnitude[i, j + 1]) and (magnitude[i, j] >= magnitude[i, j - 1]):
                    new_magnitude[i, j] = magnitude[i, j]
            elif (22.5 <= angle[i, j] < 67.5):
                if (magnitude[i, j] >= magnitude[i + 1, j + 1]) and (magnitude[i, j] >= magnitude[i - 1, j - 1]):
                    new_magnitude[i, j] = magnitude[i, j]
            elif (67.5 <= angle[i, j] < 112.5):
                if (magnitude[i, j] >= magnitude[i + 1, j]) and (magnitude[i, j] >= magnitude[i - 1, j]):
                    new_magnitude[i, j] = magnitude[i, j]
            elif (112.5 <= angle[i, j] < 157.5):
                if (magnitude[i, j] >= magnitude[i - 1, j + 1]) and (
                        magnitude[i, j] >= magnitude[i + 1, j - 1]):
                    new_magnitude[i, j] = magnitude[i, j]

    # 4. 双阈值处理，确定边缘
    strong_edges = (new_magnitude > max_threshold)
    thresholded_edges = np.zeros_like(new_magnitude)
    thresholded_edges[(new_magnitude >= min_threshold) & (new_magnitude <= max_threshold)] = 1
    weak_edges = (thresholded_edges > 0) & (strong_edges == 0)

    # 5. 边缘连接
    edge_result = np.zeros_like(new_magnitude)
    for i in range(1, height - 1):
        for j in range(1, width - 1):
            if strong_edges[i, j]:
                edge_result[i, j] = 1
            elif weak_edges[i, j]:
                if ((edge_result[i + 1, j] == 1) or (edge_result[i - 1, j] == 1)
                        or (edge_result[i, j + 1] == 1) or (edge_result[i, j - 1] == 1)
                        or (edge_result[i + 1, j + 1] == 1) or (edge_result[i - 1, j - 1] == 1)
                        or (edge_result[i - 1, j + 1] == 1) or (edge_result[i + 1, j - 1] == 1)):
                    edge_result[i, j] = 1

    return edge_result


# 打开视频文件
cap = cv2.VideoCapture("pigplayball.mp4")

while True:
    ret, frame = cap.read()
    if ret == False:
        break

    # 将图像灰度化
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

    # 调用自己实现的 canny() 函数进行边缘检测
    edge_result = canny(gray, 50, 150)

    # 显示边缘检测结果
    cv2.imshow("Canny Edge Detection Result", edge_result * 255)

    # 按键’q’退出
    if cv2.waitKey(1) == ord("q"):
        break

# 释放资源
cap.release()
cv2.destroyAllWindows()
