import cv2
import random
import numpy as np
from scipy.ndimage import label, distance_transform_edt

########################### 1. 超参数 ###########################
Annotate_CLASSES = ('背景','术中超声','吻合器','乳胶管','推结器','肝带','钳夹','超声刀','脂肪','双极电凝','棉球','血管阻断夹','肿瘤','针','线','韧带','胆囊静脉','肝脏','胆囊','分离钳','止血海绵','肝总管','胆总管','吸引器','剪刀','止血纱布','生物夹','无损伤钳','喷洒','胆囊管','胆囊动脉','电凝','标本袋','引流管','纱布','金属钛夹') # 待分类的类
Annotate_PALETTE = [[0,0,0],[138,251,213],[136,162,196],[197,83,181],[202,202,200],[113,102,140],[66,115,82],[240,16,116],[155,132,0],[155,62,0],[146,175,236],[255,172,159],[245,161,0],[134,124,118], [0,157,142], [181,85,105], [42,8,66], [255,91,0],[255,234,0],[85, 107, 179],[181, 227, 14],[72, 0, 255],[0, 155, 33],[255,0,255],[29, 32, 136],[160, 15, 95],[0,160,233],[52,184,178],[90,120,41],[255,0,0],[177,0,0],[167,24,233],[112,113,150],[0,255,0],[255,255,255],[0,255,255]]


def skeletonize(image):
    """骨架化函数，确保线条连通性并缩减为1像素宽"""
    skeleton = np.zeros_like(image)
    temp_image = np.copy(image)
    
    while True:
        eroded = cv2.erode(temp_image, None)  # 腐蚀操作
        temp_dilate = cv2.dilate(eroded, None)  # 膨胀操作
        temp = cv2.subtract(temp_image, temp_dilate)  # 提取边缘
        skeleton = cv2.bitwise_or(skeleton, temp)  # 将边缘加入骨架
        temp_image = np.copy(eroded)
        if cv2.countNonZero(temp_image) == 0:
            break
    return skeleton

# 1. *** 边缘检测并膨胀 ***
def edge_detection(image):
    """对图像的各个通道进行边缘检测并进行膨胀处理"""
    b_channel, g_channel, r_channel = cv2.split(image)
    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    edges_b = cv2.Canny(b_channel, 100, 200)
    edges_g = cv2.Canny(g_channel, 100, 200)
    edges_r = cv2.Canny(r_channel, 100, 200)
    edges_gray = cv2.Canny(gray_image, 100, 200)

    # 合并所有边缘检测结果
    edges = cv2.bitwise_or(edges_b, edges_g)
    edges = cv2.bitwise_or(edges, edges_r)
    edges = cv2.bitwise_or(edges, edges_gray)

    # 创建膨胀核并进行膨胀操作
    kernel = np.ones((3, 3), np.uint8)
    dilated_image = cv2.dilate(edges, kernel, iterations=1)

    return dilated_image

# 2. ** 检测连通区域 ***
def detect_connected_regions(dilated_image):
    """检测图像中的连通区域并过滤掉小区域"""
    _, binary = cv2.threshold(dilated_image, 1, 255, cv2.THRESH_BINARY_INV)
    binary[binary > 0] = 1  # 转换为二值图像

    # 标记连通区域
    structure = np.ones((3, 3), dtype=int)
    labeled_array, num_features = label(binary, structure=structure)
    
    # 清除掉小于100像素的区域
    filtered_labeled_array = np.copy(labeled_array)
    for label_num in range(1, num_features + 1):
        area = np.sum(labeled_array == label_num)
        if area < 100:
            filtered_labeled_array[filtered_labeled_array == label_num] = 0

    # 重新标记过滤后的连通区域
    filtered_labeled_array, num_features = label(filtered_labeled_array, structure=structure)
    return filtered_labeled_array, num_features # 返回过滤后的labeled_array和未过滤的labeled_array（用于寻找颜色）

# 3. *** 分水岭填充白色区域 ***
def fill_white_regions(filtered_labeled_array):
    """使用分水岭算法填充白色区域"""
    # 准备三通道图像作为分水岭算法的输入
    color_image = np.zeros((filtered_labeled_array.shape[0], filtered_labeled_array.shape[1], 3), dtype=np.uint8)
    
    # 将过滤后的 labeled_array 转化为 32 位整型，作为分水岭的 markers
    markers = np.copy(filtered_labeled_array).astype(np.int32)
    markers[markers == 0] = -1  # 背景标记为 -1

    # 执行分水岭算法
    cv2.watershed(color_image, markers)

    # 更新分水岭结果
    filled_labeled_array = markers.astype(int)
    filled_labeled_array[filled_labeled_array == -1] = 0

    # 使用距离变换，计算边缘像素（0）最近的非零值
    non_zero_mask = filled_labeled_array != 0
    distance, nearest_indices = distance_transform_edt(non_zero_mask == 0, return_indices=True)
    nearest_values = filled_labeled_array[tuple(nearest_indices)]
    filled_labeled_array[filled_labeled_array == 0] = nearest_values[filled_labeled_array == 0]

    return filled_labeled_array

# 4. *** 对连通区域上色(使用“filtered_labeled_array”作为颜色判断，给“filled_labeled_array”上色) ***
def color_connected_regions(filled_labeled_array, filtered_labeled_array, ori_labeled_image, Annotate_PALETTE):
    """根据原始图像的颜色和注解调色板给连通区域上色"""

    # 初始化一个三通道的彩色图像
    colored_image = np.zeros((*filled_labeled_array.shape, 3), dtype=np.uint8)

    # 遍历filtered_labeled_array中的每个标签
    unique_labels = np.unique(filtered_labeled_array)
    
    for label_num in unique_labels:
        if label_num == 0:
            continue  # 跳过背景标签
        
        # 找到filtered_labeled_array中等于当前标签的区域
        mask_filtered = (filtered_labeled_array == label_num)
        
        # 获取ori_labeled_image中对应区域的RGB值
        region_rgb_values = ori_labeled_image[mask_filtered]
        
        if len(region_rgb_values) == 0:
            continue
        
        # 计算区域的RGB平均值
        average_rgb = np.mean(region_rgb_values, axis=0)
        
        # 找到Annotate_PALETTE中与average_rgb最接近的颜色
        closest_palette_color = find_closest_palette_color(average_rgb, Annotate_PALETTE)
        
        # 将该颜色赋给filled_labeled_array对应区域的元素
        mask_filled = (filled_labeled_array == label_num)
        colored_image[mask_filled] = closest_palette_color

    return colored_image

# 寻找最近邻颜色
def find_closest_palette_color(average_rgb, Annotate_PALETTE):
    """根据平均RGB值找到Annotate_PALETTE中最接近的颜色"""
    Annotate_PALETTE = [[color[2], color[1], color[0]] for color in Annotate_PALETTE]

    average_rgb = np.array(average_rgb)
    min_distance = float('inf')
    closest_color = None
    
    # 遍历调色板，计算每个颜色与平均RGB的欧几里得距离
    for palette_color in Annotate_PALETTE:
        palette_color = np.array(palette_color)
        distance = np.linalg.norm(average_rgb - palette_color)  # 欧几里得距离
        
        if distance < min_distance:
            min_distance = distance
            closest_color = palette_color
    
    return closest_color

# 5. 对Array区域上色（4的简化版）
def Tool_color_connected_array(Array):
    colored_image = np.zeros((*Array.shape, 3), dtype=np.uint8)
    for label_num in range(1, np.max(Array) + 1):
        color = [np.random.randint(0, 254) for _ in range(3)]
        colored_image[Array == label_num] = color
    return colored_image

if __name__ == '__main__':
    """超参数"""
    image_path = './2023_02_03_09_13_48.00_08_04_21.Still085.png'
    
    """主函数，处理图像并保存结果"""
    # 读取图像
    image = cv2.imread(image_path)

    # 1. 边缘检测并膨胀
    dilated_image = edge_detection(image)
    cv2.imwrite('./1_1_range_image.png', dilated_image)

    # 2. 检测连通区域
    filtered_labeled_array, _ = detect_connected_regions(dilated_image)
    colored_image_filtered = Tool_color_connected_array(filtered_labeled_array)
    cv2.imwrite('./2_colored_image_filtered.png', colored_image_filtered)

    # 3. 分水岭填充白色区域
    filled_labeled_array = fill_white_regions(filtered_labeled_array)
    colored_image_filled = Tool_color_connected_array(filled_labeled_array)
    cv2.imwrite('./3_colored_image_filled.png', colored_image_filled)

    # 4. 对连通区域上色
    ori_labeled_image = image
    colored_image_final = color_connected_regions(filled_labeled_array, filtered_labeled_array, ori_labeled_image, Annotate_PALETTE)
    cv2.imwrite('./4_color_image_Final.png', colored_image_final)
    
    
    print("处理后的图片已保存。")
