import cv2
import os
import numpy as np

def calculate_noise_count(img_obj, w, h):
        """
        计算邻域非白色的个数
        Args:
            img_obj: img obj
            w: width
            h: height
        Returns:
            count (int)
        """
        count = 0
        width, height = img_obj.shape
        for _w_ in [w - 1, w, w + 1]:
            for _h_ in [h - 1, h, h + 1]:
                if _w_ > width - 1:
                    continue
                if _h_ > height - 1:
                    continue
                if _w_ == w and _h_ == h:
                    continue
                if img_obj[_w_, _h_] < 230:  # 二值化的图片设置为255
                    count += 1
        # print(count)
        return count
def gray(name):
    image_path = 'img_rename/'+name
    img = cv2.imread(image_path, 1)
    # 灰度
    gray_img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    k = 4
    w, h = gray_img.shape
    for _w in range(w):
        for _h in range(h):
            if _w == 0 or _h == 0:
                gray_img[_w, _h] = 255
                continue
            # 计算邻域pixel值小于255的个数
            pixel = gray_img[_w, _h]
            if pixel == 255:
                continue
            if calculate_noise_count(gray_img, _w, _h) < k:
                gray_img[_w, _h] = 255

    blur = cv2.GaussianBlur(gray_img, (9, 9), 0)  # 高斯滤波降噪
    # ● ret3代表返回的阈值。
    # ● dst代表阈值分割结果图像，与原始图像具有相同的大小和类型。
    # ● th3代表要进行阈值分割的图像，可以是多通道的，8位或32位浮点型数值。
    # ● thresh代表要设定的阈值。
    # ● maxval代表当type参数为THRESH_BINARY或者THRESH_BINARY_INV类型时，需要设定的最大值。
    # ● type代表阈值分割的类型，具体类型值如表所示。
    # ret3, th3 = cv2.threshold(blur, 134, 255, 0, gray_img)  # 二值化图片
    # ret3, th3 = cv2.threshold(blur, 134, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)  # 二值化图片
    ret3, th3 = cv2.threshold(blur, 134, 255, 0)  # 二值化图片
    # 腐蚀处理
    kernel = np.ones((3, 3), np.float32)
    img_erode = cv2.erode(th3, kernel)
    # cv2.imshow('Image', img_erode)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()
    # 创建输出文件夹
    output_folder = 'img_gray'
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    cv2.imwrite("img_gray/"+name, img_erode)

def removeBlack():
    # 读取图片
    image = cv2.imread('AIimg/1.jpg')
    # 将黑色部分替换为白色
    image[np.where((image == [0, 0, 0]).all(axis=2))] = [255, 255, 255]

    # 将红色部分保留，其余部分替换为白色
    lower_red = np.array([0, 0, 150])
    upper_red = np.array([100, 100, 255])
    mask = cv2.inRange(image, lower_red, upper_red)
    image[np.where(mask == 0)] = [255, 255, 255]

    cv2.imwrite('AIimg/1.jpg', image)
    # 显示处理后的图像
    cv2.imshow('Result', image)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

if __name__ == "__main__":
    name = "AIimg/output_image.jpg"
    # gray(name)
    removeBlack()
    # gray(name)
'''
    # 指定文件夹路径
    folder_path = 'img_rename'
    # 获取文件夹下所有文件的名字
    file_names = os.listdir(folder_path)
    # 输出文件名字
    for file_name in file_names:
        print(file_name)
        gray(file_name)
'''