import numpy as np
import cv2
import matplotlib.pyplot as plt

# 读取图片
image = cv2.imread('p41.png')

# 将图片转换为RGB格式
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

# 增强饱和度
hsv = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2HSV)
hsv[..., 1] = hsv[..., 1] * 6  # 增强饱和度
enhanced_image_rgb = cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB)

# 进行Prewitt边缘检测
gray_image = cv2.cvtColor(enhanced_image_rgb, cv2.COLOR_RGB2GRAY)
prewittx = cv2.Sobel(gray_image, cv2.CV_64F, 1, 0, ksize=3)
prewitty = cv2.Sobel(gray_image, cv2.CV_64F, 0, 1, ksize=3)
prewitt_edge = cv2.magnitude(prewittx, prewitty)

_, thresholded_image = cv2.threshold(prewitt_edge, 50, 255, cv2.THRESH_BINARY)

# 寻找轮廓
contours, hierarchy = cv2.findContours(thresholded_image.astype(np.uint8), cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE)

# 过滤面积小于阈值的轮廓
filtered_contours = []
# 计算和显示每个轮廓的最小斜率
for contour in contours:
    rect = cv2.minAreaRect(contour)  # 获取最小外接矩形
    box = cv2.boxPoints(rect)  # 计算矩形的四个角点
    box = np.intp(box)  # 角点坐标整数化  # 角点坐标整数化

    # 求边长
    widths = [np.linalg.norm(box[i] - box[(i + 1) % 4]) for i in range(4)]
    sorted_widths = sorted(widths)

    # 宽和高
    width, height = sorted_widths[0], sorted_widths[1]

    # 角度调整
    angle = rect[2]
    # if width < height:
    #     angle += 90
    area = width * height
    # 计算斜率
    slope = np.tan(np.radians(angle))
    if angle < 45 and area > 5000:
        filtered_contours.append(contour)

# 将二值图像转换为RGB格式，以便能在其上绘制彩色轮廓
contoured_image = cv2.cvtColor(thresholded_image.astype(np.uint8), cv2.COLOR_GRAY2RGB)

# 使用红色来绘制轮廓
cv2.drawContours(contoured_image, filtered_contours, -1, (255, 0, 0), 2)

# 使用matplotlib显示图像
plt.figure(figsize=(8, 8))
plt.imshow(contoured_image)
plt.title('Filtered Image with Red Contours and Slope Info')
plt.axis('off')
plt.show()