186 lines
8.1 KiB
Python
186 lines
8.1 KiB
Python
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("处理后的图片已保存。")
|