admin/Seg_Data_Server
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

first commit

Browse Source
This commit is contained in:
admin
2026-05-20 15:05:35 +08:00
commit ac09b26253
2048 changed files with 189478 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

445
Tool-可视化/0_图片Labels生成/4_deal_labels.py Normal file
View File

@@ -0,0 +1,445 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*
import os,time,sys,threading, colorsys, argparse
import asyncio, cv2, multiprocessing, random
from PIL import Image
import numpy as np
from Tool_deal_labels import edge_detection, detect_connected_regions, Tool_color_connected_array, fill_white_regions, color_connected_regions
def getFileList(dir,Filelist=[], ext=None, Max_layer=1, layer=0, Donot_Search=['1_边缘检测并膨胀', '2_连通区域检测', '3_分水岭算法填充']):
"""
获取文件夹及其子文件夹中文件列表
输入 dir：文件夹根目录
输入 ext: 扩展名
返回： 文件路径列表
"""
newDir = dir
if os.path.isfile(dir):
if ext is None:
Filelist.append(dir)
else:
if ext in dir[-3:]:
Filelist.append(dir)
elif os.path.isdir(dir):
file_name = os.path.basename(dir)
# 判断是否在禁搜名单中
if file_name in Donot_Search:
return Filelist
for s in os.listdir(dir):
newDir=os.path.join(dir,s)
if layer <= Max_layer:
getFileList(newDir, Filelist, ext, Max_layer, layer+1)
return Filelist
class Deal_image():
def __init__(self, Annotate_CLASSES = ('肝脏','胆囊'), Annotate_PALETTE = [[255,91,0],[255,234,0]], src_label_fold = "./Label", save_pro_label_fold = "./LABEL_PNG_new", save_GT_label_fold = "./Label_Generate", GT_channel = 1, pro_append_name="_label", GT_append_name="_gtFine_labelTrainIds", ori_img_folder="./ORI_PNG", res_label_folder="./Result_label", save_merge_pic_folder="./Result_merge", back_gnd_color=0, first_class_color=1, pic_type="png", Max_width = 10000, Label_Max_Search_layer=1000, save_process_pics=False, bg_PALETTE = [0,0,0]):
# 背景最好放在最后
# self.src_CLASSES = ('肝脏','胆囊','分离钳','止血海绵','肝总管','胆总管','吸引器','剪刀','止血纱布','生物夹','无损伤钳','喷洒','胆囊管','胆囊动脉','电凝','标本袋','引流管','纱布','金属钛夹','术中超声','吻合器','乳胶管','推结器','肝带','钳夹','超声刀','脂肪','双极电凝','棉球','血管阻断夹','肿瘤','针','线','韧带','胆囊静脉','背景')
# self.src_PALETTE = np.array([[255,91,0],[255,234,0],[85, 111, 181],[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],[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],[0,0,0]])
# self.src_CLASSES_NUM = np.shape(self.src_CLASSES)[0]
self.bg_PALETTE = bg_PALETTE # 背景颜色 TODO
self.Annotate_CLASSES = Annotate_CLASSES # 待分类的类
self.Annotate_PALETTE = np.array(Annotate_PALETTE) # 每一类的像素直
self.Annotate_CLASSES_NUM = np.shape(Annotate_CLASSES)[0] # 类数量
self.save_process_pics = save_process_pics # 保存中间过程图片
self.src_label_fold = src_label_fold # 原始标签图片 保存位置
self.save_pro_label_fold = save_pro_label_fold # 优化后标签图片 保存位置
self.save_GT_label_fold = save_GT_label_fold # GT标签图片 保存位置
self.ori_img_folder = ori_img_folder # 最原始手术图片 保存位置
self.res_label_folder = res_label_folder # 训练出来的label 保存位置
self.save_merge_pic_folder = save_merge_pic_folder # 融合图像保存位置
self.pro_append_name = pro_append_name # 优化后标签图片后缀
self.GT_append_name = GT_append_name # GT标签图片后缀
self.GT_channel = GT_channel # GT标签图片通道数
self.Max_width = Max_width # 最大图片宽度（匹配时候用）
self.pic_type = pic_type # 图片类型
self.back_gnd_color = back_gnd_color # 背景颜色
self.first_class_color = first_class_color # 第一类上的颜色
self.Label_Max_Search_layer=Label_Max_Search_layer # 文件夹最大搜索深度
try:
self.labellist_src = getFileList(src_label_fold, [], pic_type, self.Label_Max_Search_layer)
print('本次执行检索到ori_label图片 '+str(len(self.labellist_src))+' 张图像')
except:
self.labellist_src = None
print("没有ori_label相关文件")
try:
# print(save_pro_label_fold)
self.labellist_pro = getFileList(save_pro_label_fold, [], pic_type, self.Label_Max_Search_layer)
print('本次执行检索到pro_label图片 '+str(len(self.labellist_pro))+' 张图像')
except:
self.labellist_pro = None
print("没有pro_label相关文件")
try:
self.imglist_src = getFileList(ori_img_folder, [], pic_type, self.Label_Max_Search_layer)
self.reslist_src = getFileList(res_label_folder, [], pic_type, self.Label_Max_Search_layer)
print('本次执行检索到ori原始图片 '+str(len(self.imglist_src))+' 张图像')
print('本次执行检索到训练train_result图片 '+str(len(self.reslist_src))+' 张图像')
except:
self.imglist_src = None
self.reslist_src = None
print("没有train_result和原始图片相关文件")
# 获取单张图片各个通路信息
def get_single_pic_rgb(self, imgpath):
print(imgpath)
image = Image.open(imgpath).convert('RGB') # 转为RGB图片
# 将 RGB 色值分离
image.load()
r, g, b = image.split()
r = np.array(r)
g = np.array(g)
b = np.array(b)
return image, r, g, b
# 将单个pro图片变成GT图片
def Conver_pro_label_pic_2_GT_pic(self, imgpath, imgname):
time_start=time.time() # 记录开始时间
# 获取单张图片各个通路信息
image, r,g,b = self.get_single_pic_rgb(imgpath)
result_gt = np.ones(np.shape(image))*self.back_gnd_color # 初始化填充内容为back_gnd_color
gt_number = self.first_class_color # 第一类上色颜色确定
# PALETTE中排除掉 '背景' [0,0,0]
PALETTE_No_Bg = self.Annotate_PALETTE[~np.all(self.Annotate_PALETTE == self.bg_PALETTE, axis=1)]
# 遍历所有待识别颜色
for [Annotate_PALETTE_r, Annotate_PALETTE_g, Annotate_PALETTE_b] in PALETTE_No_Bg:
# 查找三原色匹配位置
locate_r = np.where( r == Annotate_PALETTE_r )
locate_g = np.where( g == Annotate_PALETTE_g )
locate_b = np.where( b == Annotate_PALETTE_b )
# 查找都匹配位置（交集）
# 将矩阵换一种表示形式
locate_r = np.array(locate_r[0]) * self.Max_width + np.array(locate_r[1])
locate_g = np.array(locate_g[0]) * self.Max_width + np.array(locate_g[1])
locate_b = np.array(locate_b[0]) * self.Max_width + np.array(locate_b[1])
# 用自带函数寻找匹配项
matched = np.intersect1d(np.intersect1d(locate_r, locate_g), locate_b)
matched = np.concatenate(([matched // self.Max_width], [np.mod(matched, self.Max_width)]), 0)
result_gt[matched[0],matched[1], :] = gt_number
gt_number = gt_number + 1
# 输出GT图片
if(int(self.GT_channel) == 1):
result_gt = result_gt[:,:,0]
elif(int(self.GT_channel) == 3):
result_gt = cv2.cvtColor(np.float32(result_gt), cv2.COLOR_RGB2BGR) # rgb颜色互换
else:
print("GT_channel 必须为1或3")
quit
try: # 新建文件夹
os.mkdir(self.save_GT_label_fold)
except:
print("已有"+self.save_GT_label_fold)
if imgname.lower().endswith(('.jpg', '.png')):
save_dir = os.path.join(self.save_GT_label_fold, os.path.basename(imgname).rpartition('.')[0]+self.GT_append_name+'.'+self.pic_type)
else:
save_dir = os.path.join(self.save_GT_label_fold, os.path.basename(imgname)+self.GT_append_name+'.'+self.pic_type)
cv2.imwrite(save_dir, result_gt)
print("GT图片已保存", save_dir)
time_end=time.time() # 输出结束时间
print('time cost',time_end-time_start,'s')
# 将处理好的图片转化为GT图片
def Conver_pro_label_pic_2_GT_pic_all(self):
print("\033[33m**** 进行转换将Pro_label_pic转换为GT_label_pic ****\033[0m")
print("\033[33mPro_label_pic存储位置为：\033[0m", self.save_pro_label_fold)
print("\033[33mGT_label_pic生成位置为：\033[0m", self.save_GT_label_fold)
try:
# print(save_pro_label_fold)
self.labellist_pro = getFileList(save_pro_label_fold, [], pic_type, self.Label_Max_Search_layer)
print('本次执行检索到pro_label图片 '+str(len(self.labellist_pro))+' 张图像')
except:
self.labellist_pro = None
print("没有pro_label相关文件")
try:
os.mkdir(self.save_GT_label_fold) # 新建存储文件夹
except:
print("已有"+self.save_GT_label_fold)
# 指定最大进程数为 3
max_processes = 20
# 创建Pool对象
pool = multiprocessing.Pool(processes=max_processes)
# 创建并启动进程
args_list1 = []
args_list2 = []
# 遍历整个文件夹
for imgpath in self.labellist_pro:
imgname = os.path.basename(imgpath).rpartition('.')[0].replace(self.pro_append_name,"")
args_list1.append(imgpath)
args_list2.append(imgname)
args_list = zip(args_list1, args_list2)
# 使用进程池并行执行任务
pool.starmap(self.Conver_pro_label_pic_2_GT_pic, args_list)
# 关闭进程池
pool.close()
pool.join()
def Conver_ori_label_pic_2_pro_pic(self, imgpath, imgname):
time_start=time.time() # 记录开始时间
# 获取单张图片各个通路信息
image = cv2.imread(imgpath)
# 1. 边缘检测并膨胀
dilated_image = edge_detection(image)
# 如果需要存储中间态图片
if(self.save_process_pics == True):
if imgname.lower().endswith(('.jpg', '.png')):
save_dir = os.path.join(self.save_pro_label_fold, '1_边缘检测并膨胀', os.path.basename(imgname).rpartition('.')[0]+self.pro_append_name+'_Edge'+'.'+self.pic_type)
else:
save_dir = os.path.join(self.save_pro_label_fold, '1_边缘检测并膨胀', os.path.basename(imgname)+self.pro_append_name+'_Edge'+'.'+self.pic_type)
cv2.imwrite(save_dir, dilated_image)
print("中间态-边缘检测并膨胀 图片已保存", save_dir)
time_end=time.time() # 输出结束时间
print('time cost',time_end-time_start,'s')
# 2. 检测连通区域
filtered_labeled_array, _ = detect_connected_regions(dilated_image)
colored_image_filtered = Tool_color_connected_array(filtered_labeled_array)
# 如果需要存储中间态图片
if(self.save_process_pics == True):
if imgname.lower().endswith(('.jpg', '.png')):
save_dir = os.path.join(self.save_pro_label_fold, '2_连通区域检测', os.path.basename(imgname).rpartition('.')[0]+self.pro_append_name+'_Region'+'.'+self.pic_type)
else:
save_dir = os.path.join(self.save_pro_label_fold, '2_连通区域检测', os.path.basename(imgname)+self.pro_append_name+'_Region'+'.'+self.pic_type)
cv2.imwrite(save_dir, colored_image_filtered)
print("中间态-连通区域检测 图片已保存", save_dir)
time_end=time.time() # 输出结束时间
print('time cost',time_end-time_start,'s')
# 3. 分水岭填充白色区域
filled_labeled_array = fill_white_regions(filtered_labeled_array)
colored_image_filled = Tool_color_connected_array(filled_labeled_array)
# 如果需要存储中间态图片
if(self.save_process_pics == True):
if imgname.lower().endswith(('.jpg', '.png')):
save_dir = os.path.join(self.save_pro_label_fold, '3_分水岭算法填充', os.path.basename(imgname).rpartition('.')[0]+self.pro_append_name+'_FillEdge'+'.'+self.pic_type)
else:
save_dir = os.path.join(self.save_pro_label_fold, '3_分水岭算法填充', os.path.basename(imgname)+self.pro_append_name+'_FillEdge'+'.'+self.pic_type)
cv2.imwrite(save_dir, colored_image_filled)
print("中间态-分水岭算法填充 图片已保存", save_dir)
time_end=time.time() # 输出结束时间
print('time cost',time_end-time_start,'s')
# 4. 对连通区域最终上色
ori_labeled_image = image
result_pro = color_connected_regions(filled_labeled_array, filtered_labeled_array, ori_labeled_image, self.Annotate_PALETTE)
if imgname.lower().endswith(('.jpg', '.png')):
save_dir = os.path.join(self.save_pro_label_fold, os.path.basename(imgname).rpartition('.')[0]+self.pro_append_name+'.'+self.pic_type)
else:
save_dir = os.path.join(self.save_pro_label_fold, os.path.basename(imgname)+self.pro_append_name+'.'+self.pic_type)
print("Pro图片已保存", save_dir)
cv2.imwrite(save_dir, result_pro)
time_end=time.time() # 输出结束时间
print('time cost',time_end-time_start,'s')
# 将原始src图片转化为处理好的pro图片
def Conver_ori_label_pic_2_pro_pic_all(self):
print("\033[33m**** 进行转换将Ori_label_pic转换为Pro_label_pic ****\033[0m")
print("\033[33mOri_label_pic存储位置为：\033[0m", self.src_label_fold)
print("\033[33mPro_label_pic生成位置为：\033[0m", self.save_pro_label_fold)
# 输出颜色预处理图片
try:
os.mkdir(self.save_pro_label_fold) # 新建存储文件夹
except:
print("已有"+self.save_pro_label_fold)
if(self.save_process_pics == True):
try:
os.mkdir(os.path.join(self.save_pro_label_fold, '1_边缘检测并膨胀')) # 新建存储1_边缘检测并膨胀文件夹
except:
print("已有"+os.path.join(self.save_pro_label_fold, '1_边缘检测并膨胀'))
try:
os.mkdir(os.path.join(self.save_pro_label_fold, '2_连通区域检测')) # 新建存储2_连通区域检测文件夹
except:
print("已有"+os.path.join(self.save_pro_label_fold, '2_连通区域检测'))
try:
os.mkdir(os.path.join(self.save_pro_label_fold, '3_分水岭算法填充')) # 新建存储1_边缘检测并膨胀文件夹
except:
print("已有"+os.path.join(self.save_pro_label_fold, '3_分水岭算法填充'))
# 指定最大进程数为 20，多参数函数并行
max_processes = 20
# 创建Pool对象
pool = multiprocessing.Pool(processes=max_processes)
# 创建并启动进程
args_list1 = []
args_list2 = []
# 遍历整个文件夹
for imgpath in self.labellist_src:
if imgpath.lower().endswith(('.jpg', '.png')):
imgname= os.path.basename(imgpath).rpartition('.')[0].replace(self.pro_append_name,"")
else:
imgname= os.path.basename(imgpath).replace(self.pro_append_name,"")
try:
print("Processing: ", imgname, "...")
# self.Conver_ori_label_pic_2_pro_pic(imgpath, imgname)s
# args_list.append({'imgpath': imgpath, 'imgname': imgname})
args_list1.append(imgpath)
args_list2.append(imgname)
except:
os.system("echo "+imgname+" >> error_1.txt")
args_list = zip(args_list1, args_list2)
# 使用进程池并行执行任务
pool.starmap(self.Conver_ori_label_pic_2_pro_pic, args_list) # 使用starmap进行多参数并行
# 关闭进程池
pool.close()
pool.join()
# 图片堆叠
def Merge_ori_pic_and_label_pic(self, res_img_path, res_imgname):
time_start=time.time() # 记录开始时间
# 获取单张图片各个通路信息
ori_img_path = os.path.join(self.ori_img_folder, res_imgname+'.'+self.pic_type)
if not os.path.exists(ori_img_path):
print("****照片不存在：****", ori_img_path)
return -1
ori_image, ori_r, ori_g, ori_b = self.get_single_pic_rgb(ori_img_path)
res_image, res_r, res_g, res_b = self.get_single_pic_rgb(res_img_path)
merge_img = np.array(ori_image) # merge图片初始化，默认图片背景为0.0.0
# 遍历所有待识别颜色
for [Annotate_PALETTE_r, Annotate_PALETTE_g, Annotate_PALETTE_b] in self.Annotate_PALETTE:
# 查找三原色匹配位置
locate_r = np.where( res_r == Annotate_PALETTE_r )
locate_g = np.where( res_g == Annotate_PALETTE_g )
locate_b = np.where( res_b == Annotate_PALETTE_b )
# 查找都匹配位置（交集）
# 将矩阵换一种表示形式
locate_r = np.array(locate_r[0]) * self.Max_width + np.array(locate_r[1])
locate_g = np.array(locate_g[0]) * self.Max_width + np.array(locate_g[1])
locate_b = np.array(locate_b[0]) * self.Max_width + np.array(locate_b[1])
# 用自带函数寻找匹配项
matched = np.intersect1d(np.intersect1d(locate_r, locate_g), locate_b)
matched = np.concatenate(([matched // self.Max_width], [np.mod(matched, self.Max_width)]), 0)
merge_img[matched[0],matched[1], 0] = Annotate_PALETTE_r
merge_img[matched[0],matched[1], 1] = Annotate_PALETTE_g
merge_img[matched[0],matched[1], 2] = Annotate_PALETTE_b
# 转成cv2形式
merge_img = cv2.cvtColor(np.float32(merge_img), cv2.COLOR_RGB2BGR)
try: # 新建文件夹
os.mkdir(self.save_merge_pic_folder)
except:
print("已有"+self.save_merge_pic_folder)
if res_imgname.lower().endswith(('.jpg', '.png')):
save_dir = os.path.join(self.save_merge_pic_folder, os.path.basename(res_imgname).rpartition('.')[0]+'.'+self.pic_type)
else:
save_dir = os.path.join(self.save_merge_pic_folder, os.path.basename(res_imgname)+'.'+self.pic_type)
cv2.imwrite(save_dir, merge_img)
print("Merge图片已保存", save_dir)
time_end=time.time() # 输出结束时间
print('time cost',time_end-time_start,'s')
# 将label图片与原图片重合
def Merge_ori_pic_and_label_pic_all(self):
# 遍历整个文件夹
for res_img_path in self.reslist_src:
if res_img_path.lower().endswith(('.jpg', '.png')):
res_imgname = os.path.basename(res_img_path).rpartition('.')[0].replace(self.pro_append_name,"")
else:
res_imgname = os.path.basename(res_img_path).replace(self.pro_append_name,"")
print("Processing: ", res_imgname, "...")
self.Merge_ori_pic_and_label_pic(res_img_path, res_imgname)
if __name__ == "__main__":
Annotate_CLASSES = ('肝脏','胆囊','分离钳','止血海绵','肝总管','胆总管','吸引器','剪刀','止血纱布','生物夹','无损伤钳','喷洒','胆囊管','胆囊动脉','电凝','标本袋','引流管','纱布','金属钛夹','术中超声','吻合器','乳胶管','推结器','肝带','钳夹','超声刀','脂肪','双极电凝','棉球','血管阻断夹','肿瘤','针','线','韧带','胆囊静脉','背景') # 待分类的类
Annotate_PALETTE = [[255,91,0],[255,234,0],[85, 111, 181],[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],[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],[0,0,0]] # 每一类的像素直
bg_PALETTE = [0,0,0] # 背景的RGB
# 创建参数解析器
parser = argparse.ArgumentParser(description='Process some files.')
# 添加参数选项
parser.add_argument('-src_fold', dest='src_label_fold', default='./', help='source label folder')
parser.add_argument('-save_pro_fold', dest='save_pro_label_fold', default='./save_pro_label_fold', help='processed label folder')
parser.add_argument('-save_GT_fold', dest='save_GT_label_fold', default='./save_GT_label_fold', help='ground truth folder')
parser.add_argument('-fold_search_depth', dest='Label_Max_Search_layer', default='1000', type=int, help='Folder Search Depth')
parser.add_argument('-pro_suffix_name', dest='pro_append_name', default='_label', help='Pro file suffix')
parser.add_argument('-GT_suffix_name', dest='GT_append_name', default='_gtFine_labelTrainIds', help='GT file suffix')
parser.add_argument('-GT_channel', dest='GT_channel', default='1', type=int, help='GT file channel(1 or 3)')
parser.add_argument('-back_gnd_color', dest='back_gnd_color', default='0', type=int, help='Color of "Back ground"(0 or 255)')
parser.add_argument('-first_class_color', dest='first_class_color', default='1', type=int, help='Color of "First Class"')
parser.add_argument('-pic_type', dest='pic_type', default='png', help='type of picture(Do not add ".")')
parser.add_argument('-Max_width', dest='Max_width', default='10000', type=int, help='Max width of picture')
parser.add_argument('-Rebuild_from', dest='Rebuild_from', default='label', help='Source to Rebuild Labels(label/pro)')
parser.add_argument('-Rebuild_to', dest='Rebuild_to', default='GT', help='Destination of Rebuild Labels(pro/GT)')
parser.add_argument('-save_process_pics', dest='save_process_pics', default='False', help='Save the processed pics(e.g.Gray_pics,Color_pics) in generating pro_pics')
# 解析命令行参数
args = parser.parse_args()
src_label_fold = args.src_label_fold
save_pro_label_fold = args.save_pro_label_fold
save_GT_label_fold = args.save_GT_label_fold
Label_Max_Search_layer = args.Label_Max_Search_layer
pro_append_name = args.pro_append_name
GT_append_name = args.GT_append_name
GT_channel = args.GT_channel
back_gnd_color = args.back_gnd_color
first_class_color = args.first_class_color
pic_type = args.pic_type
Max_width = args.Max_width
Rebuild_from = args.Rebuild_from
Rebuild_to = args.Rebuild_to
save_process_pics = args.save_process_pics
try: # 遍历文件深度，最小为1
Label_Max_Search_layer=int(Label_Max_Search_layer)
except:
Label_Max_Search_layer=1000
try: # GT标签图片通道数
GT_channel=int(GT_channel)
except:
GT_channel=1
try: # 背景颜色（背景选择0或255)
back_gnd_color=int(back_gnd_color)
except:
back_gnd_color=0
try: # 第一类上的颜色(如果背景为0,选择1；)
first_class_color=int(first_class_color)
except:
first_class_color=1
try: # 最大图片宽度（匹配时候用）
Max_width=int(Max_width)
except:
Max_width=10000
if(save_process_pics.lower() == 'false'):
save_process_pics = False
elif(save_process_pics.lower() == 'true'):
save_process_pics = True
else:
save_process_pics = False
D = Deal_image(Annotate_CLASSES=Annotate_CLASSES, Annotate_PALETTE=Annotate_PALETTE, src_label_fold=src_label_fold, save_pro_label_fold=save_pro_label_fold, save_GT_label_fold=save_GT_label_fold, GT_channel=GT_channel, pro_append_name=pro_append_name, GT_append_name=GT_append_name, back_gnd_color=back_gnd_color, first_class_color=first_class_color, pic_type=pic_type, Max_width=Max_width, Label_Max_Search_layer=Label_Max_Search_layer, save_process_pics=save_process_pics, bg_PALETTE = bg_PALETTE)
# print(D.src_CLASSES_NUM)
if Rebuild_from == 'label':
# 1.先将所有原始图片转为pro图片
D.Conver_ori_label_pic_2_pro_pic_all()
pass
if Rebuild_to == 'GT':
# 2.再将pro图片转为GT图片
D.Conver_pro_label_pic_2_GT_pic_all()
pass

185
Tool-可视化/0_图片Labels生成/Tool_deal_labels.py Normal file
View File

@@ -0,0 +1,185 @@
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("处理后的图片已保存。")

93
Tool-可视化/Tool_Check_and_Gen_Txt_Label_ori_label.py Normal file
View File

@@ -0,0 +1,93 @@
import os
import cv2
import numpy as np
from collections import Counter
from PIL import Image
# ※ 需要修改输入输出路径 ※ #
input_dir = 'Data/ann_dir'
output_dir_1 = 'Data/labels/train'
# output_dir_2 = 'Data/labels/val'
output_dir_2 = 'Data/labels/train'
os.makedirs(output_dir_1, exist_ok=True)
os.makedirs(output_dir_2, exist_ok=True)
# 全局统计颜色频率
global_color_counter = Counter()
global_class_counter = Counter()
# 自动提取颜色映射（跳过背景）
def extract_color_mapping(img_path):
img = Image.open(img_path).convert('RGB')
pixels = list(img.getdata())
counter = Counter(pixels)
color_map = {}
for color, count in counter.items():
global_color_counter[color] += count
if color != (0, 0, 0) and color[0] == color[1] == color[2]:
class_id = color[0] - 1
if class_id >= 0:
color_map[color] = class_id
global_class_counter[class_id] += count
return color_map
# 处理单张图片
def process_image(img_path, save_path_list):
color_to_class = extract_color_mapping(img_path)
if not color_to_class:
print(f"[跳过] {os.path.basename(img_path)} 无有效目标")
return
img = cv2.imread(img_path)
h, w = img.shape[:2]
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
lines = []
for rgb, class_id in color_to_class.items():
mask = np.all(img_rgb == rgb, axis=-1).astype(np.uint8) * 255
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
for contour in contours:
if contour.shape[0] < 3:
continue
norm_pts = contour.squeeze(1).astype(np.float32)
norm_pts[:, 0] /= w
norm_pts[:, 1] /= h
flat = norm_pts.flatten()
line = f"{class_id} " + " ".join([f"{x:.6f}" for x in flat])
lines.append(line)
if lines:
for save_path in save_path_list:
with open(save_path, 'w') as f:
for line in lines:
f.write(line + "\n")
print(f"[✔] 转换成功: {os.path.basename(save_path)}，共 {len(lines)} 个实例")
else:
print(f"[⚠] {os.path.basename(img_path)} 没有轮廓")
# 主执行逻辑：批量处理
for fname in os.listdir(input_dir):
if fname.lower().endswith('.png'):
img_path = os.path.join(input_dir, fname)
base_name = os.path.splitext(fname)[0]
txt_path_1 = os.path.join(output_dir_1, base_name + '.txt')
txt_path_2 = os.path.join(output_dir_2, base_name + '.txt')
process_image(img_path, [txt_path_1, txt_path_2])
# 打印颜色统计和类别映射
print("\n📊 所有图像颜色统计：")
for color, count in global_color_counter.most_common():
if color == (0, 0, 0):
print(f"背景颜色 {color} 出现次数: {count}")
elif color[0] == color[1] == color[2]:
class_id = color[0] - 1
print(f"颜色 {color} → class {class_id}，出现次数: {count}")
else:
print(f"[⚠] 非灰阶颜色 {color}，跳过")
print("\n✅ 有效类别统计（class_id → 总像素数）：")
for class_id in sorted(global_class_counter):
print(f"class {class_id}: {global_class_counter[class_id]} pixels")
print(f"\n✅ 全部图像处理完毕。标签输出目录：{output_dir_1}、{output_dir_2}")

109
Tool-可视化/Tool_Check_and_Gen_Txt_Label_sort_label.py Normal file
View File

@@ -0,0 +1,109 @@
import os
import cv2
import numpy as np
from collections import Counter, defaultdict
from PIL import Image
# ※ 需要修改输入输出路径 ※ #
input_dir = 'Data/ann_dir'
output_dir_1 = 'Data/labels/train'
output_dir_2 = 'Data/labels/train'
os.makedirs(output_dir_1, exist_ok=True)
os.makedirs(output_dir_2, exist_ok=True)
# 全局统计颜色频率与 class 像素频率
global_color_counter = Counter()
global_class_counter = Counter()
color_class_counter = defaultdict(int) # (R,G,B) → count
color_to_old_class = {} # (R,G,B) → class_id
remap_class_dict = {} # old_class_id → new_class_id
# 自动提取颜色映射（跳过背景）
def extract_color_mapping(img_path):
img = Image.open(img_path).convert('RGB')
pixels = list(img.getdata())
counter = Counter(pixels)
color_map = {}
for color, count in counter.items():
global_color_counter[color] += count
if color != (0, 0, 0) and color[0] == color[1] == color[2]:
class_id = color[0] - 1
if class_id >= 0:
color_map[color] = class_id
global_class_counter[class_id] += count
color_class_counter[color] += count
color_to_old_class[color] = class_id
return color_map
# 处理单张图片
def process_image(img_path, save_path_list):
color_to_class = extract_color_mapping(img_path)
if not color_to_class:
print(f"[跳过] {os.path.basename(img_path)} 无有效目标")
return
img = cv2.imread(img_path)
h, w = img.shape[:2]
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
lines = []
for rgb, old_class_id in color_to_class.items():
mask = np.all(img_rgb == rgb, axis=-1).astype(np.uint8) * 255
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
for contour in contours:
if contour.shape[0] < 3:
continue
norm_pts = contour.squeeze(1).astype(np.float32)
norm_pts[:, 0] /= w
norm_pts[:, 1] /= h
flat = norm_pts.flatten()
new_class_id = remap_class_dict.get(old_class_id, old_class_id)
line = f"{new_class_id} " + " ".join([f"{x:.6f}" for x in flat])
lines.append(line)
if lines:
for save_path in save_path_list:
with open(save_path, 'w') as f:
for line in lines:
f.write(line + "\n")
print(f"[✔] 转换成功: {os.path.basename(save_path)}，共 {len(lines)} 个实例")
else:
print(f"[⚠] {os.path.basename(img_path)} 没有轮廓")
# 第一次遍历图像，仅提取颜色信息用于构建 class 映射
for fname in os.listdir(input_dir):
if fname.lower().endswith('.png'):
img_path = os.path.join(input_dir, fname)
extract_color_mapping(img_path)
# 构建 class_id 重映射表（按像素数从大到小排序）
class_pixel_count = {color_to_old_class[c]: count for c, count in color_class_counter.items()}
sorted_classes = sorted(class_pixel_count.items(), key=lambda x: x[1], reverse=True)
remap_class_dict = {old_cls: new_idx for new_idx, (old_cls, _) in enumerate(sorted_classes)}
# 第二次遍历图像，正式处理并生成标签
for fname in os.listdir(input_dir):
if fname.lower().endswith('.png'):
img_path = os.path.join(input_dir, fname)
base_name = os.path.splitext(fname)[0]
txt_path_1 = os.path.join(output_dir_1, base_name + '.txt')
txt_path_2 = os.path.join(output_dir_2, base_name + '.txt')
process_image(img_path, [txt_path_1, txt_path_2])
# 打印颜色统计和类别映射
print("\n📊 所有图像颜色统计：")
for color, count in global_color_counter.most_common():
if color == (0, 0, 0):
print(f"背景颜色 {color} 出现次数: {count}")
elif color[0] == color[1] == color[2]:
class_id = color[0] - 1
print(f"颜色 {color} → class {class_id}，出现次数: {count}")
else:
print(f"[⚠] 非灰阶颜色 {color}，跳过")
print("\n✅ 有效类别统计（原 class_id → 新 class_id → 总像素数）：")
for old_id, new_id in remap_class_dict.items():
print(f"class {old_id} → {new_id}: {class_pixel_count[old_id]} pixels")
print(f"\n✅ 全部图像处理完毕。标签输出目录：{output_dir_1}、{output_dir_2}")

25
Tool-可视化/Tool_Gen_8_Bit_PNG[没用，不认].py Normal file
View File

@@ -0,0 +1,25 @@
import os
from PIL import Image
import numpy as np
input_dir = 'Data/ann_dir' # 输入标签图所在目录
output_dir = 'Data/labels/train' # YOLO 格式输出目录
os.makedirs(output_dir, exist_ok=True)
for filename in os.listdir(input_dir):
if not filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp')):
continue
input_path = os.path.join(input_dir, filename)
output_path = os.path.join(output_dir, os.path.splitext(filename)[0] + '.png')
# 打开图像并转为 RGB，再提取 R 通道（等价于 G、B）
img_rgb = Image.open(input_path).convert('RGB')
r_channel = img_rgb.split()[0] # 或使用 .getchannel('R')
# 保存为 8-bit 单通道 PNG
r_channel.save(output_path, format='PNG')
print(f"Saved: {output_path}")
print("✅ 所有标签图已转换为 8-bit 单通道灰度图。")

92
Tool-可视化/get_FPS.py Normal file
View File

@@ -0,0 +1,92 @@
import warnings
warnings.filterwarnings('ignore')
import argparse
import logging
import math
import os
import random
import time
import sys
from copy import deepcopy
from pathlib import Path
from threading import Thread
import numpy as np
import torch.distributed as dist
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torch.optim.lr_scheduler as lr_scheduler
import torch.utils.data
import yaml
from torch.cuda import amp
from torch.nn.parallel import DistributedDataParallel as DDP
from tqdm import tqdm
import os
os.environ["HTTP_PROXY"] = "http://127.0.0.1:2089"
os.environ["HTTPS_PROXY"] = "http://127.0.0.1:2089"
from ultralytics import YOLO
from ultralytics.utils.torch_utils import select_device
from ultralytics.nn.tasks import attempt_load_weights
def get_weight_size(path):
stats = os.stat(path)
return f'{stats.st_size / 1024 / 1024:.1f}'
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--weights', type=str, default='yolov8n.pt', help='trained weights path')
parser.add_argument('--batch', type=int, default=1, help='total batch size for all GPUs')
parser.add_argument('--imgs', nargs='+', type=int, default=[640, 640], help='[height, width] image sizes')
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--warmup', default=200, type=int, help='warmup time')
parser.add_argument('--testtime', default=1000, type=int, help='test time')
parser.add_argument('--half', action='store_true', default=False, help='fp16 mode.')
opt = parser.parse_args()
device = select_device(opt.device, batch=opt.batch)
# Model
weights = opt.weights
if weights.endswith('.pt'):
model = attempt_load_weights(weights, device=device, fuse=True)
print(f'Loaded {weights}') # report
else:
model = YOLO(weights).model
model.fuse()
model = model.to(device)
example_inputs = torch.randn((opt.batch, 3, *opt.imgs)).to(device)
if opt.half:
model = model.half()
example_inputs = example_inputs.half()
print('begin warmup...')
for i in tqdm(range(opt.warmup), desc='warmup....'):
model(example_inputs)
print('begin test latency...')
time_arr = []
for i in tqdm(range(opt.testtime), desc='test latency....'):
if device.type == 'cuda':
torch.cuda.synchronize()
start_time = time.time()
model(example_inputs)
if device.type == 'cuda':
torch.cuda.synchronize()
end_time = time.time()
time_arr.append(end_time - start_time)
std_time = np.std(time_arr)
infer_time_per_image = np.sum(time_arr) / (opt.testtime * opt.batch)
if weights.endswith('.pt'):
print(f'model weights:{opt.weights} size:{get_weight_size(opt.weights)}M (bs:{opt.batch})Latency:{infer_time_per_image:.5f}s +- {std_time:.5f}s fps:{1 / infer_time_per_image:.1f}')
else:
print(f'model yaml:{opt.weights} (bs:{opt.batch})Latency:{infer_time_per_image:.5f}s +- {std_time:.5f}s fps:{1 / infer_time_per_image:.1f}')

31
Tool-可视化/inference.py Normal file
View File

@@ -0,0 +1,31 @@
import os
from ultralytics import YOLO
from pathlib import Path
# === 设置路径 ===
weights_path = './runs/segment/train6/weights/best.pt' # 模型权重
source_dir = './Data/images/train' # 输入图片目录
output_dir = './Data/result/train' # 推理结果输出目录
os.makedirs(output_dir, exist_ok=True)
# === 加载模型 ===
model = YOLO(weights_path)
# === 推理参数 ===
results = model.predict(
source=source_dir, # 可为单图像路径、目录、视频、摄像头索引等
save=True, # 是否保存图像
save_txt=False, # 是否保存标签（分割任务不常用）
save_conf=True, # 是否保存置信度
save_crop=False, # 是否保存目标裁剪图
project=output_dir, # 输出根路径
name='', # 子文件夹名（为空即直接输出到 output_dir）
exist_ok=True, # 允许覆盖已有文件夹
imgsz=640, # 推理分辨率（默认640）
conf=0.02, # 置信度阈值 # 设置的越高，分类越少
iou=0.45, # NMS IOU 阈值
# device='cuda:0' # 改为 'cpu' 如果没有GPU
device='cpu' # ✅ 修改为 CPU 模式
)
print(f"✅ 推理完成，结果保存在：{output_dir}")

28
Tool-可视化/my_dataset.yaml Normal file
View File

@@ -0,0 +1,28 @@
# my_dataset.yaml
path: ./Data # 数据集根路径
train: images/train # 训练集图片目录
val: images/train # 如果没有验证集，可暂用训练集代替
# ※ label放置位置 ※
# 放在 labels 目录下，格式为：
# - labels/train/*.txt
# 或者还放在images/train/*.txt中
# 针对8bit PNG格式
# 类别名，按你的实际颜色映射定义（可继续补充）
names:
0: class_0
1: class_1
2: class_2
3: class_3
4: class_4
# # 针对TXT格式
# names:
# 0: background # 默认第0类是背景
# 1: class_0
# 2: class_1
# 分割任务专用字段
# 如果你的 mask 是单通道 PNG（如 labelTrainIds.png），确保它们是 8-bit 灰度图
# Ultralytics YOLO 会自动根据 names 映射类别

7
Tool-可视化/train.py Normal file
View File

@@ -0,0 +1,7 @@
from ultralytics import YOLO
# Load a pretrained YOLO11 segment model
model = YOLO("yolo11n-seg.pt")
# Train the model
results = model.train(data="./my_dataset.yaml", epochs=100, imgsz=640)

139
Tool-可视化/yolov11_heatmap_V1.py Normal file
View File

@@ -0,0 +1,139 @@
import os
import cv2
import numpy as np
import torch
from ultralytics import YOLO
from pytorch_grad_cam import GradCAM, GradCAMPlusPlus, XGradCAM, EigenCAM, HiResCAM, LayerCAM, RandomCAM, EigenGradCAM, KPCA_CAM
from pytorch_grad_cam.utils.image import show_cam_on_image
from pytorch_grad_cam.base_cam import BaseCAM
class ActivationMaximizationTarget:
def __init__(self, channel=0):
self.channel = channel
def __call__(self, model_output):
# model_output: [B, C, H, W]
return model_output[:, self.channel, :, :].mean()
# 1. 中间封装类：只返回 [B, C, H, W] 特征图
class YoloFeatureExtractor(torch.nn.Module):
def __init__(self, model):
super().__init__()
self.model = model
def forward(self, x):
out = self.model(x)
if isinstance(out, (list, tuple)) and len(out) > 1:
feat_candidates = out[1]
if isinstance(feat_candidates, (list, tuple)):
for i, feat in enumerate(feat_candidates):
if isinstance(feat, torch.Tensor) and feat.ndim == 4:
print(f"[DEBUG] 找到特征图: item[1][{i}] -> shape={feat.shape}")
feat.requires_grad_() # 🔥 关键修复点
return feat
else:
print(f"[DEBUG] item[1][{i}] 类型: {type(feat)}")
raise RuntimeError("未找到可用于 CAM 的特征图")
# 2. 加载 YOLOv11 模型
# model_path = r"runs\segment\train2\weights\best.pt"
model_path = "yolo11n-seg.pt" # 替换为你的模型路径
model = YOLO(model_path)
model.model.eval()
# 3. 提取 CAM hook 层（最后一个 Conv2d）
target_layers = []
for module in model.model.modules():
if isinstance(module, torch.nn.Conv2d):
target_layers.append(module)
if not target_layers:
raise RuntimeError("未找到卷积层")
target_layers = [target_layers[-1]] # 使用最后一层 Conv2d
# 4. 输出文件夹初始化
output_root = "result_CAM_Method"
os.makedirs(output_root, exist_ok=True)
cam_methods = {
"GradCAM": GradCAM,
"GradCAMPlusPlus": GradCAMPlusPlus,
"XGradCAM": XGradCAM,
"EigenCAM": EigenCAM,
"HiResCAM": HiResCAM,
"LayerCAM": LayerCAM,
"RandomCAM": RandomCAM,
"EigenGradCAM": EigenGradCAM,
"KPCA_CAM": KPCA_CAM
}
for method_name in cam_methods:
os.makedirs(os.path.join(output_root, method_name), exist_ok=True)
# 5. 遍历图像
input_dir = "Data/img_dir"
for img_name in os.listdir(input_dir):
if img_name.lower().endswith((".jpg", ".jpeg", ".png", ".bmp")):
img_path = os.path.join(input_dir, img_name)
orig_image = cv2.imread(img_path)
if orig_image is None:
continue
orig_image = cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB)
orig_h, orig_w = orig_image.shape[:2]
# letterbox resize + padding to 640x640
target_size = 640
scale = min(target_size / orig_w, target_size / orig_h)
new_w = int(orig_w * scale)
new_h = int(orig_h * scale)
resized_image = cv2.resize(orig_image, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
pad_w = target_size - new_w
pad_h = target_size - new_h
pad_left = pad_w // 2
pad_right = pad_w - pad_left
pad_top = pad_h // 2
pad_bottom = pad_h - pad_top
pad_color = (114, 114, 114)
padded_image = cv2.copyMakeBorder(resized_image, pad_top, pad_bottom, pad_left, pad_right,
cv2.BORDER_CONSTANT, value=pad_color)
# 图像归一化 + tensor 转换
padded_image_float = padded_image.astype(np.float32) / 255.0
device = next(model.model.parameters()).device
input_tensor = torch.from_numpy(padded_image_float.transpose(2, 0, 1))[None].to(device)
# 6. 用 wrapper 包装模型以兼容 CAM
wrapped_model = YoloFeatureExtractor(model.model)
# 遍历每种 CAM 方法
for method_name, cam_class in cam_methods.items():
with cam_class(model=wrapped_model, target_layers=target_layers) as cam:
target = [ActivationMaximizationTarget(channel=0)]
cam_result = cam(input_tensor=input_tensor, targets=target)[0]
# 如果输出为 3 维（如 [C, H, W]），取通道平均为 [H, W]
if isinstance(cam_result, torch.Tensor):
cam_result = cam_result.detach().cpu().numpy()
if cam_result.ndim == 3:
cam_result = cam_result.mean(axis=0)
elif cam_result.ndim != 2:
raise ValueError(f"[CAM ERROR] Unexpected CAM shape: {cam_result.shape}")
grayscale_cam = cam_result # 安全赋值
cam_cropped = grayscale_cam[pad_top:pad_top + new_h, pad_left:pad_left + new_w]
if cam_cropped.size == 0:
continue
cam_resized = cv2.resize(cam_cropped, (orig_w, orig_h), interpolation=cv2.INTER_LINEAR)
cam_resized = cam_resized - cam_resized.min()
cam_resized = cam_resized / (cam_resized.max() + 1e-8)
orig_image_float = orig_image.astype(np.float32) / 255.0
overlay_image = show_cam_on_image(orig_image_float, cam_resized, use_rgb=True)
heatmap = cv2.applyColorMap(np.uint8(255 * cam_resized), cv2.COLORMAP_JET)
heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
base_name = os.path.splitext(img_name)[0]
overlay_path = os.path.join(output_root, method_name, f"{base_name}_overlay.jpg")
heatmap_path = os.path.join(output_root, method_name, f"{base_name}_heatmap.jpg")
cv2.imwrite(overlay_path, cv2.cvtColor(overlay_image, cv2.COLOR_RGB2BGR))
cv2.imwrite(heatmap_path, cv2.cvtColor(heatmap, cv2.COLOR_RGB2BGR))

189
Tool-可视化/yolov11_heatmap_V2.py Normal file
View File

@@ -0,0 +1,189 @@
import os
import cv2
import numpy as np
import torch
from ultralytics import YOLO
from pytorch_grad_cam import (
GradCAM, GradCAMPlusPlus, XGradCAM, EigenCAM,
HiResCAM, LayerCAM, RandomCAM, EigenGradCAM, KPCA_CAM
)
from pytorch_grad_cam.utils.image import show_cam_on_image
from pytorch_grad_cam.base_cam import BaseCAM
class ActivationMaximizationTarget:
def __init__(self, channel=0):
self.channel = channel
def __call__(self, model_output):
if model_output.ndim == 4:
# [B, C, H, W]
return model_output[:, self.channel, :, :].mean()
elif model_output.ndim == 3:
# [C, H, W]
return model_output[self.channel, :, :].mean()
else:
raise ValueError(f"Unsupported model_output shape: {model_output.shape}")
class YoloFeatureExtractor(torch.nn.Module):
def __init__(self, model):
super().__init__()
self.model = model
def forward(self, x):
out = self.model(x)
if isinstance(out, (list, tuple)) and len(out) > 1:
if isinstance(out, (list, tuple)):
feat_candidates = out[1] if len(out) > 1 else out[0]
if isinstance(feat_candidates, (list, tuple)):
for i, feat in enumerate(feat_candidates):
if isinstance(feat, torch.Tensor) and feat.ndim == 4:
feat.requires_grad_()
return feat
raise RuntimeError("未找到可用于 CAM 的特征图")
# ------------------------------ 设置路径 -------------------------------
# model_path = "yolo11n-seg.pt" # ← 替换为你的模型路径
model_path = r"runs\segment\train6\weights\best.pt"
model_name = os.path.splitext(os.path.basename(model_path))[0]
output_root = f"result_CAM_Method_{model_name}"
input_dir = r"Data\images\train"
# ------------------------------ 加载模型 -------------------------------
model = YOLO(model_path)
model.model.eval()
device = next(model.model.parameters()).device
# ------------------------------ 提取所有卷积层 -------------------------------
conv_layers = []
for idx, layer in enumerate(model.model.modules()):
if isinstance(layer, torch.nn.Conv2d):
conv_layers.append((idx, layer))
# conv_layers.reverse()
if not conv_layers:
raise RuntimeError("未找到卷积层")
# ------------------------------ 设置 CAM 方法 -------------------------------
cam_methods = {
"GradCAM": GradCAM,
"GradCAMPlusPlus": GradCAMPlusPlus,
"XGradCAM": XGradCAM,
# "EigenCAM": EigenCAM, # 风险较高，容易炸内存
"HiResCAM": HiResCAM,
"LayerCAM": LayerCAM,
"RandomCAM": RandomCAM,
"EigenGradCAM": EigenGradCAM,
# "KPCA_CAM": KPCA_CAM # 风险较高，容易炸内存
}
# 创建目录
for method in cam_methods:
os.makedirs(os.path.join(output_root, method), exist_ok=True)
# ------------------------------ 处理图像 -------------------------------
for img_name in os.listdir(input_dir):
if not img_name.lower().endswith((".jpg", ".jpeg", ".png", ".bmp")):
continue
img_path = os.path.join(input_dir, img_name)
orig_image = cv2.imread(img_path)
if orig_image is None:
continue
orig_image = cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB)
orig_h, orig_w = orig_image.shape[:2]
# Resize + padding
target_size = 640
scale = min(target_size / orig_w, target_size / orig_h)
new_w, new_h = int(orig_w * scale), int(orig_h * scale)
resized_image = cv2.resize(orig_image, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
pad_w, pad_h = target_size - new_w, target_size - new_h
pad_left, pad_right = pad_w // 2, pad_w - pad_w // 2
pad_top, pad_bottom = pad_h // 2, pad_h - pad_h // 2
padded_image = cv2.copyMakeBorder(resized_image, pad_top, pad_bottom, pad_left, pad_right,
cv2.BORDER_CONSTANT, value=(114, 114, 114))
padded_image_float = padded_image.astype(np.float32) / 255.0
input_tensor = torch.from_numpy(padded_image_float.transpose(2, 0, 1))[None].to(device)
input_tensor.requires_grad_() # ★★★ 加这行！
wrapped_model = YoloFeatureExtractor(model.model)
# -------------------- 遍历每一层 + 每个方法 ---------------------
for layer_idx, layer in conv_layers:
print(f"\nProcessing Layer {layer_idx}: {layer.__class__.__name__}")
for method_name, cam_class in cam_methods.items():
try:
# 方法执行前预检查特征图尺寸
if cam_class in [EigenCAM, KPCA_CAM]:
# 特征图太大提前跳过
try:
# 临时 forward 一次拿特征图大小
with torch.no_grad():
feat = layer(input_tensor)
feat_shape = feat.shape # [B, C, H, W]
numel = feat.numel()
if numel > 4096 ** 2: # 超过 16M 元素就跳过
print(f"[SKIP] {method_name} on Layer {layer_idx}: 特征图过大 shape={feat_shape}")
continue
except Exception as e:
print(f"[SKIP] {method_name} on Layer {layer_idx}: 特征图检查失败: {e}")
continue
print(f" Using {method_name}...")
# Way 1: 使用 wrapper 包装模型
# with cam_class(model=wrapped_model, target_layers=[layer]) as cam:
# targets = [ActivationMaximizationTarget(channel=0)]
# Way 2:
with cam_class(model=model.model, target_layers=[layer]) as cam:
targets = [ActivationMaximizationTarget(channel=0)]
cam_output = cam(input_tensor=input_tensor, targets=targets)
# 正确顺序：先判断类型，再使用变量
if isinstance(cam_output, (list, tuple)):
cam_result = cam_output[0]
else:
cam_result = cam_output
# Tensor → Numpy
if isinstance(cam_result, torch.Tensor):
cam_result = cam_result.detach().cpu().numpy()
# 处理不同维度
if cam_result.ndim == 4:
cam_result = cam_result[0].mean(axis=0)
elif cam_result.ndim == 3:
cam_result = cam_result.mean(axis=0)
elif cam_result.ndim == 2:
pass # OK
else:
print(f"[SKIP] {method_name} on Layer {layer_idx}：CAM 结果维度异常 {cam_result.shape}")
continue
# EigenCAM 特征图过大保护
if cam_class in [EigenCAM, KPCA_CAM] and cam_result.size > 4096**2:
print(f"[SKIP] {method_name} on Layer {layer_idx} 特征图太大，跳过")
continue
# CAM 后处理
cam_cropped = cam_result[pad_top:pad_top + new_h, pad_left:pad_left + new_w]
cam_resized = cv2.resize(cam_cropped, (orig_w, orig_h), interpolation=cv2.INTER_LINEAR)
cam_resized = (cam_resized - cam_resized.min()) / (cam_resized.max() + 1e-8)
overlay_image = show_cam_on_image(orig_image.astype(np.float32) / 255.0,
cam_resized, use_rgb=True)
heatmap = cv2.applyColorMap(np.uint8(255 * cam_resized), cv2.COLORMAP_JET)
heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
layer_name = layer.__class__.__name__
base = os.path.splitext(img_name)[0]
fname = f"{layer_idx}_{layer_name}_{base}"
overlay_path = os.path.join(output_root, method_name, f"{fname}_overlay.jpg")
heatmap_path = os.path.join(output_root, method_name, f"{fname}_heatmap.jpg")
cv2.imwrite(overlay_path, cv2.cvtColor(overlay_image, cv2.COLOR_RGB2BGR))
cv2.imwrite(heatmap_path, cv2.cvtColor(heatmap, cv2.COLOR_RGB2BGR))
except Exception as e:
print(f"[ERROR] {method_name} on Layer {layer_idx} failed: {e}")

7
Tool-可视化/使用流程.txt Normal file
View File

@@ -0,0 +1,7 @@
0. cd 0_图片Labels生成； python 4_deal_labels.py 对label进行处理；# 放入Data\ann_dir中
1. 【推荐】使用 Tool_Check_and_Gen_Txt_Label_sort_label.py 生成数据【按照从大到小顺序0~N个class】
使用 Tool_Check_and_Gen_Txt_Label_ori_label.py 生成数据【原始图片中的class】
2. 修改 my_dataset.yaml 修改路径、分类【分类数和步骤1程序生成数量相同】；
2.1. 删除label下的.cache文件
3. python train.py 、 python inference.py 【修改pt模型】 进行训练、推理
4. python yolov11_heatmap_V2.py 进行热图可视化