# -*- coding: utf-8 -*-
# @Time    : 2024/4/15 9:23
# @Author  : len
# @File    : 感兴趣区域提取.py
# @Software: PyCharm 
# @Comment :
import os

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

plt.rcParams['font.sans-serif'] = ['SimHei']

# 函数：加载图片，找到屏幕轮廓并显示
def process_image(image_path):
    # 加载图片
    image = cv2.imread(image_path)

    # 转换图片为灰度图
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # 应用边缘检测
    edges = cv2.Canny(gray, 50, 150)

    # 寻找轮廓
    contours, _ = cv2.findContours(edges, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

    # 按面积大小排序轮廓，并取最大的10个
    contours = sorted(contours, key=cv2.contourArea, reverse=True)[:10]

    screen_contour = None
    for contour in contours:
        # 近似轮廓
        perimeter = cv2.arcLength(contour, True)
        approximation = cv2.approxPolyDP(contour, 0.02 * perimeter, True)

        # 屏幕轮廓应该有4个点（矩形）
        if len(approximation) == 4:
            screen_contour = approximation
            break

    if screen_contour is not None:
        # 在原图上绘制轮廓
        cv2.drawContours(image, [screen_contour], -1, (0, 255, 0), 3)
        screen_image = cv2.polylines(image, [screen_contour], True, (0,255,0), 3)
    else:
        screen_image = image

    # 展示结果
    plt.imshow(cv2.cvtColor(screen_image, cv2.COLOR_BGR2RGB))
    plt.title('屏幕检测')
    plt.show()

    return screen_image

# 设置文件夹路径
folder_path = './data/pingmu'  # 将此路径替换为你的图片文件夹路径

# 遍历文件夹中的所有文件
for filename in os.listdir(folder_path):
    if filename.lower().endswith(('.png', '.jpg', '.jpeg')):  # 检查文件扩展名
        file_path = os.path.join(folder_path, filename)  # 获取文件完整路径
        process_image(file_path)  # 处理图片