dice_test.py

import numpy as np
from tqdm import tqdm
from MedSAM.segment_anything import sam_model_registry
from MedSAM.demo import BboxPromptDemo
import os
import argparse
from CAM.main_vit import show_mask, dice_coeff
from Utils import *
import logging
import time
from CAM.utils import scoremap2bbox
from skimage import measure
import multiprocessing



def scale_cam_image(cam, target_size=None):
    result = []
    for img in cam:
        img = img - np.min(img)
        img = img / (1e-7 + np.max(img))
        if target_size is not None:
            img = cv2.resize(img, target_size)
        result.append(img)
    result = np.float32(result)

    return result


def preprossess(img, mask):
    for i in range(img.shape[0]):
        for j in range(img.shape[1]):
            if (img[i][j] <= 1e-9):
                mask[i][j] = 0
    return mask


def find_max_eras(mask):
    # 图像读取
    img = mask
    img = np.array(img)
    img[img != 0] = 1  # 图像二值化
    # 图像实例化
    img = measure.label(img, connectivity=2)
    props = measure.regionprops(img)
    # 最大区域获取
    max_area = 0
    max_index = 0
    # props只包含像素值不为零区域的属性，因此index要从1开始
    for index, prop in enumerate(props, start=1):
        if prop.area > max_area:
            max_area = prop.area
            # index 代表每个联通区域内的像素值；prop.area代表相应连通区域内的像素个数
            max_index = index
    if max_index == 0:
        return img
    img[img != max_index] = 0
    img[img == max_index] = 1

    return img


def show_mask_image(mask, ax, random_color=False, alpha=0.95):
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([alpha])], axis=0)
    else:
        color = np.array([251 / 255, 252 / 255, 30 / 255, alpha])
    h, w = mask.shape[-2:]
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    ax.imshow(mask_image)


def main(opt, logger):
    n_jobs = multiprocessing.cpu_count()
    device = "cuda:0"
    # CAM模型导入
    # 初始化dice
    dice_caa = 0
    dice_sam_caa = 0
    dice_out = 0
    dice_in = 0
    pic_out_len = 0
    pic_in_len = 0
    # 保存全部图片的灰度平均值
    mean_intensity = 0
    out_intensity = 0
    in_intensity = 0
    pic_intensity_in=0
    pic_intensity_out = 0
    # 获得视频列表
    video_path = build_video_list(opt.video_dir)
    cnt = 0

    for video_dir in tqdm(video_path, desc="Processing Videos"):  # 遍历所以的视频

        video_name = video_dir.split('\\')[-1]
        video_output_path = os.path.join(opt.video_out_path, video_dir.split('/')[-1])
        if not os.path.exists(video_output_path):
            os.makedirs(video_output_path)
        cam_list = []
        label_res_list = []
        cam_box_list = []
        SAM_list = []
        cam_box_dice = []
        cam_dice = []
        sam2_dice = []
        label_dice = []
        #保存SAM_stack
        SAM_stack=[]
        # 保存全部的box框
        box_list = []
        # 保存原始cam热力图
        cam_ori_mask = []
        cam_stack_mask = np.zeros((240, 240))
        # 保存处理后的cam热力图
        cam_pro_mask = []
        cam_pro_stack_mask = np.zeros((240, 240))
        #保存sam处理图
        sam_pro_mask=[]
        # 保存box圈主的区域
        box_mask = []
        box_stack_mask = np.zeros((240, 240))
        # 保存label叠加
        label_mask = []
        label_stack_mask = np.zeros((240, 240))
        # 得到视频全部图片帧
        frame_names, frame = build_frame_list(video_dir)
        cam_path = opt.cam_dir
        label_path = opt.label_path
        # 得到视频全部图片的label
        label = build_label_list(label_path, frame_names)

        # 得到视频全部图片的cam热力图
        for img_name, label_data in zip(frame_names, label):
            cnt += 1
            img_data = np.load(os.path.join(opt.img_path, img_name.split('.')[0] + '.npz'))["arr_0"]
            grad_cam = np.load(os.path.join(opt.out_path, img_name.split('.')[0] + '.npz'))['original_cam']
            caa_grad_cam = np.load(os.path.join(opt.out_path, img_name.split('.')[0] + '.npz'))['caa_cam']
            sam_mask = np.load(os.path.join("/root/data1/brats/box_stack_sam1_stack/twosam_addcam_1012",img_name.split('.')[0] + '.npz'))['caa_sam1_stack_pred']
            sam_pro_mask.append(sam_mask)
            # 保存原始的cam_mask
            cam_ori_mask.append(caa_grad_cam.copy())
            temp = []
            for _ in range(3):
                temp.append(img_data)
            tmpimage = np.array(temp)
            tmpimage = tmpimage.transpose(1, 2, 0)
            original_grad_cam = np.array(caa_grad_cam).copy()
            original_grad_cam[original_grad_cam < 0.8] = 0
            original_grad_cam[grad_cam < 0.2] = 0
            caa_grad_cam = preprossess(img_data, caa_grad_cam)

            caa_grad_cam[caa_grad_cam < 0.8*caa_grad_cam.max()] = 0
            caa_grad_cam[caa_grad_cam >= 0.8*caa_grad_cam.max()] = 1
            caa_grad_cam[grad_cam < 0.2] = 0

            caa_grad_cam = find_max_eras(caa_grad_cam)
            # 保存处理后的cam_mask
            cam_pro_mask.append(caa_grad_cam.copy())

            boxcaa, cntcaa = scoremap2bbox(scoremap=caa_grad_cam, threshold=0.4, multi_contour_eval=True)
            # 保存box圈主的区域
            zero_array = np.zeros_like(caa_grad_cam)
            # 使用box的坐标，圈住的区域内赋值为1
            x1, y1, x2, y2 = boxcaa[0]
            tempbox = []
            tempbox.append(boxcaa[0])
            # box_list.append(boxcaa[0])
            zero_array[y1:y2 + 1, x1:x2 + 1] = 1
            for j_ in range(1, cntcaa):
                x1, y1, x2, y2 = boxcaa[j_]
                tempbox.append(boxcaa[j_])
                zero_array[y1:y2 + 1, x1:x2 + 1] = 1
            box_list.append(np.array(tempbox))
            box_mask.append(zero_array)
            label_mask.append(label_data)


        rescnt = -1
        intensity_list=[]
        out_intensity_dangqian=0
        in_intensity_dangqian=0
        pic_intensity_out_dangqian=0
        pic_intensity_in_dangqian=0
        label_intensity=[]
        # 初始化区间字典
        intensity_bins = {
            '0.9-1.0': [],
            '0.8-0.9': [],
            '0.7-0.8': [],
            '0.6-0.7': [],
            '0.5-0.6': [],
            '0.4-0.5': [],
            '0.3-0.4': [],
            '0.2-0.3': [],
            '0.1-0.2': [],
            '0.0-0.1': []
        }
        for img_name, label_data in zip(frame_names, label):
            rescnt += 1
            cam_dice.append(0)
            img_data = np.load(os.path.join(opt.img_path, img_name.split('.')[0] + '.npz'))["arr_0"]
            grad_cam = np.load(os.path.join(opt.out_path, img_name.split('.')[0] + '.npz'))['original_cam']
            caa_grad_cam = np.load(os.path.join(opt.out_path, img_name.split('.')[0] + '.npz'))['caa_cam']
            temp = []
            for _ in range(3):
                temp.append(img_data)
            tmpimage = np.array(temp)
            tmpimage = tmpimage.transpose(1, 2, 0)
            cam_list.append(show_cam_on_image(tmpimage, caa_grad_cam, returncam=True))

            original_grad_cam = np.array(caa_grad_cam).copy()

            original_grad_cam[original_grad_cam < 0.8] = 0
            original_grad_cam[grad_cam < 0.2] = 0
            caa_grad_cam = preprossess(img_data, caa_grad_cam)
            caa_grad_cam[caa_grad_cam < 0.8] = 0
            caa_grad_cam[caa_grad_cam >= 0.8] = 1
            caa_grad_cam[grad_cam < 0.2] = 0

            caa_grad_cam = find_max_eras(caa_grad_cam)
            zero_array = np.zeros_like(caa_grad_cam)
            for j_ in box_list[rescnt]:
                boxcaa = j_
                x1, y1, x2, y2 = boxcaa
                zero_array[y1:y2 + 1, x1:x2 + 1] = 1
            caa_grad_cam[zero_array == 0] = 0

            dice = dice_coeff(caa_grad_cam, label_data)
            dice_caa += dice
            cam_box_list.append(show_cam_on_image(tmpimage.copy(), caa_grad_cam, box=box_list[rescnt], returncam=True))
            cam_box_dice.append(dice_coeff(caa_grad_cam, label_data))
            label_res_list.append(show_cam_on_image(tmpimage.copy(), label_data, returncam=True))
            label_dice.append(1)
            SAM_mask_caa = np.load(os.path.join("/root/data1/brats/box_stack_sam1_stack/twosam_addcam_1012",img_name.split('.')[0] + '.npz'))['caa_sam1_stack_pred']

            dice_sam = dice_coeff(SAM_mask_caa, label_data)
            SAM_stack.append(SAM_mask_caa)
            dice_sam_caa += dice_sam
            SAM_list.append(show_cam_on_image(tmpimage.copy(), SAM_mask_caa, returncam=True))
            sam2_dice.append(dice_sam)

            intensity=process_image(img_data,SAM_mask_caa)
            label_intensity.append(process_image(img_data,label_data))
            mean_intensity+=intensity
            intensity_list.append(intensity)

            if dice_coeff(caa_grad_cam, label_data) >= 0.1:
                dice_out += dice_sam
                pic_out_len += 1
            else:
                dice_in += dice_sam
                pic_in_len += 1

            if dice_coeff(SAM_mask_caa, label_data) >= 0.3:
                out_intensity_dangqian+=intensity
                out_intensity+=intensity
                pic_intensity_out += 1
                pic_intensity_out_dangqian += 1
            else:
                in_intensity+=intensity
                in_intensity_dangqian += intensity
                pic_intensity_in += 1
                pic_intensity_in_dangqian += 1

            if pic_intensity_out_dangqian==0:
                pic_intensity_out_dangqian=1
            if pic_intensity_in_dangqian==0:
                pic_intensity_in_dangqian=1
            # 根据Dice值将灰度值存储到对应的区间
            if 0.9 <= dice_sam <= 1.0:
                intensity_bins['0.9-1.0'].append(intensity)
            elif 0.8 <= dice_sam < 0.9:
                intensity_bins['0.8-0.9'].append(intensity)
            elif 0.7 <= dice_sam < 0.8:
                intensity_bins['0.7-0.8'].append(intensity)
            elif 0.6 <= dice_sam < 0.7:
                intensity_bins['0.6-0.7'].append(intensity)
            elif 0.5 <= dice_sam < 0.6:
                intensity_bins['0.5-0.6'].append(intensity)
            elif 0.4 <= dice_sam < 0.5:
                intensity_bins['0.4-0.5'].append(intensity)
            elif 0.3 <= dice_sam < 0.4:
                intensity_bins['0.3-0.4'].append(intensity)
            elif 0.2 <= dice_sam < 0.3:
                intensity_bins['0.2-0.3'].append(intensity)
            elif 0.1 <= dice_sam < 0.2:
                intensity_bins['0.1-0.2'].append(intensity)
            elif 0.0 <= dice_sam < 0.1:
                intensity_bins['0.0-0.1'].append(intensity)

        # 计算每个区间的平均灰度值
        avg_intensities = {}
        for dice_range, intensities in intensity_bins.items():
            if intensities:  # 确保列表不为空
                avg_intensities[dice_range] = np.mean(intensities)
            else:
                avg_intensities[dice_range] = 0  # 如果该区间没有数据，设为0

        avg_intensities["label"]=np.array(label_intensity).mean()

        # 绘制直方图
        plt.bar(avg_intensities.keys(), avg_intensities.values())
        plt.xlabel('Dice Coefficient Range')
        plt.ylabel('Average Intensity')
        plt.title('Average Intensity per Dice Coefficient Range')
        plt.xticks(rotation=45)
        # 保存直方图到指定路径，格式可以是png, jpg等
        plt.savefig(f"/root/data1/cam-MEDSAM/huidumean/{video_dir.split('/')[-1]}huidumean.png", dpi=300, bbox_inches='tight')  # 指定保存路径

        # 清空当前绘图，避免影响后续绘图
        plt.clf()

        logging.info(f"{video_name}:"
                     f"dice_sam:{round(np.array(sam2_dice).sum() / sam2_dice.__len__(), 4)},"
                     f"intensity_sam:{round(np.array(intensity_list).sum() / intensity_list.__len__(), 4)},"
                     f"dice_sam_sum: {round(dice_sam_caa / cnt, 4)},"
                     f"intensity_sam_sum: {round(mean_intensity / cnt, 4)},"
                     f"去掉0.3的intensity_sum：{round(out_intensity / pic_intensity_out, 4)},"
                     f"只看0.3以下的intensity_sum：{round(in_intensity / pic_intensity_in, 4)},"
                     f"只看0.3以下的该视频intensity：{round(out_intensity_dangqian / pic_intensity_out_dangqian, 4)},"
                     f"只看0.3以下的该视频intensity：{round(in_intensity_dangqian / pic_intensity_in_dangqian, 4)},"
                     f"数据集{pic_intensity_out}/{cnt}")


def loadLogger(args):
    logger = logging.getLogger()
    logger.setLevel(logging.INFO)
    formatter = logging.Formatter(fmt="[ %(asctime)s ] %(message)s",
                                  datefmt="%a %b %d %H:%M:%S %Y")

    sHandler = logging.StreamHandler()
    sHandler.setFormatter(formatter)

    logger.addHandler(sHandler)

    if not args.not_save:
        work_dir = os.path.join(args.work_dir,
                                time.strftime("%Y.%m.%dT%H %M %S", time.localtime()))
        if not os.path.exists(work_dir):
            os.makedirs(work_dir)

        fHandler = logging.FileHandler(work_dir + '/log.txt', mode='w')
        fHandler.setLevel(logging.DEBUG)
        fHandler.setFormatter(formatter)

        logger.addHandler(fHandler)

    return logger


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--device', default='cuda:0', help='device id (i.e. 0 or 0,1 or cpu)')
    parser.add_argument('--num_classes', type=int, default=2)
    parser.add_argument('--img_path', default="/root/data1/brats/val/yes")
    parser.add_argument('--label_path', default="/root/data1/brats/val/label")
    parser.add_argument('--cam_dir', default="/root/data1/brats/val/cam")
    parser.add_argument('--video_dir', default="/root/data1/brats/val/valyes")
    parser.add_argument('--not-save', default=False, action='store_true',
                        help='If yes, only output log to terminal.')
    parser.add_argument('--work-dir', default='./work_dir',
                        help='the work folder for storing results')
    parser.add_argument('--out_path', default="/root/data1/brats/val/cam")
    parser.add_argument('--video_out_path', default="/root/data1/brats/video_medsam_twosam_addcam_1012")
    parser.add_argument('--sam1-dir', default="/root/data1/brats/box_stack_sam1")
    parser.add_argument('--npz-dir', default="/root/data1/brats/twosam_addcam_1012")
    opt = parser.parse_args()
    logger = loadLogger(opt)
    main(opt, logger)