conda create -n medsam python=3.10 -y

conda activate medsam

pip install torch-2.0.0+cu117-cp310-cp310-win_amd64.whl -i https://pypi.tuna.tsinghua.edu.cn/simple

pip install torchvision-0.15.1+cu117-cp310-cp310-win_amd64.whl -i https://pypi.tuna.tsinghua.edu.cn/simple

pip install -e . -i https://pypi.tuna.tsinghua.edu.cn/simple

# set up the parser
parser = argparse.ArgumentParser(description='preprocess CT images')
parser.add_argument('-i', '--nii_path', type=str, default='data/FLARE22Train/images', help='path to the nii images')
parser.add_argument('-gt', '--gt_path', type=str, default='data/FLARE22Train/labels', help='path to the ground truth',)
parser.add_argument('-o', '--npz_path', type=str, default='data/Npz_files', help='path to save the npz files')
 
parser.add_argument('--image_size', type=int, default=256, help='image size')
parser.add_argument('--modality', type=str, default='CT', help='modality')
parser.add_argument('--anatomy', type=str, default='Abd-Gallbladder', help='anatomy')
parser.add_argument('--img_name_suffix', type=str, default='_0000.nii.gz', help='image name suffix')
parser.add_argument('--label_id', type=int, default=9, help='label id')
parser.add_argument('--prefix', type=str, default='CT_Abd-Gallbladder_', help='prefix')
parser.add_argument('--model_type', type=str, default='vit_b', help='model type')
parser.add_argument('--checkpoint', type=str, default='work_dir/SAM/sam_vit_b_01ec64.pth', help='checkpoint')
parser.add_argument('--device', type=str, default='cuda:0', help='device')
# seed
parser.add_argument('--seed', type=int, default=2023, help='random seed')
args = parser.parse_args()

pip install chardet
pip install --force-reinstall charset-normalizer==3.1.0

# 获取 args.gt_path 目录下所有文件名，并按字典序排序，将结果赋值给 names
names = sorted(os.listdir(args.gt_path))
# 将 args.npz_path 和 args.model_type 拼接成一个新路径名 save_path
save_path = args.npz_path + '_' + args.model_type
# 创建新的目录save_path。如果该目录已经存在，则不做任何操作。如果不存在，则新建该目录及其所有上级目录
os.makedirs(save_path, exist_ok=True)
# 打印输出 names 列表的长度即图片数量
print('image number:', len(names))

# 初始化模型，设置好 args.model_type、args.checkpoint、args.device
sam_model = sam_model_registry[args.model_type](checkpoint=args.checkpoint).to(args.device)

imgs = []  # 图片（images）
gts =  []  # 标签（labels）
img_embeddings = []  # 图片嵌入信息
 
# 遍历 ground truth 文件夹, names 默认为 args.gt_path（'data/MedSAMDemo_2D/train/labels'）下排序好的文件列表
for gt_name in tqdm(names):
    image_name = gt_name.split('.')[0] + args.img_name_suffix  # 获得文件名称 + 后缀名
    gt_data = io.imread(join(args.gt_path, gt_name))  # 获得 ground truth 数据
    if len(gt_data.shape)==3:  # 如果 gt_data 是三维的，则只取第一个通道（灰度图）
        gt_data = gt_data[:,:,0]
    assert len(gt_data.shape)==2, 'ground truth should be 2D'  # 确保分割标签是二维的（高和宽）
    # 这行代码的作用是将分割标签（即 gt_data）缩放到指定大小（args.image_size），并将其值转换为二进制形式（0 和 1）。
    # 具体来说，它会将分割标签中所有等于 args.label_id 的像素点设置为 1，其余像素点设置为 0，然后将结果缩放到指定大小。
    # 这里使用了 scikit-image 库的 transform.resize 函数，
    # 并指定了 order=0 表示使用最近邻插值法，
    # preserve_range=True 表示保持输入图像的范围不变（即值仍在 0 到 1 之间），
    # mode='constant' 表示在缩放后填充常数值的方式为使用边界值填充。
    # 最终得到的结果是一个二值图像，即只包含 0 和 1 两种像素值的图像。
    gt_data = transform.resize(gt_data==args.label_id, (args.image_size, args.image_size), order=0, preserve_range=True, mode='constant')
    # 将gt_data值转换为 8 位无符号整数
    gt_data = np.uint8(gt_data)
 
    # exclude tiny objects 如果分割标签中包含的像素点数大于 100，则执行以下操作(对源图像进行预处理，加入最终的数据集中)
    if np.sum(gt_data)>100:  
        # 最大值是 1，就两种像素点。确保分割标签是二值图
        assert np.max(gt_data)==1 and np.unique(gt_data).shape[0]==2, 'ground truth should be binary'
        # 获得图像数据
        image_data = io.imread(join(args.img_path, image_name))
        # 如果图像包含透明度通道，则只取前三个通道，即 RGB 通道
        if image_data.shape[-1]>3 and len(image_data.shape)==3:
            image_data = image_data[:,:,:3]
        # 如果图像只有一个通道，则将其复制三次，即得到一个 RGB 图像
        if len(image_data.shape)==2:
            image_data = np.repeat(image_data[:,:,None], 3, axis=-1)
        # nii preprocess start
        # 计算图像的亮度范围，即确定合适的像素值下限和上限
        # 使用 np.percentile 函数分别计算了图像中像素值从小到大排列后第 0.5% 和第 99.5% 的值，
        # 将其作为下限和上限，用于后续的像素值标准化处理
        lower_bound, upper_bound = np.percentile(image_data, 0.5), np.percentile(image_data, 99.5)
        # 将图像中的像素值限制在 lower_bound 和 upper_bound 之间
        image_data_pre = np.clip(image_data, lower_bound, upper_bound)
        # 将调整后的图像进行标准化，方法是先将图像中所有像素值减去最小值，然后除以像素值范围（即最大值减去最小值），最后乘以255，使像素值缩放到0-255的范围。
        # 这样做的目的是为了使得图像不受亮度范围的影响，并且方便后续模型的处理，因为很多模型输入都需要归一化的图像数据
        image_data_pre = (image_data_pre - np.min(image_data_pre))/(np.max(image_data_pre)-np.min(image_data_pre))*255.0
        # 将背景像素（黑色）设置为 0
        image_data_pre[image_data==0] = 0
        # 调整图像大小，并使用三次样条插值方法进行重采样，使得图像更加平滑，并保持图像的范围不变
        image_data_pre = transform.resize(image_data_pre, (args.image_size, args.image_size), order=3, preserve_range=True, mode='constant', anti_aliasing=True)
        # 将图像像素值转换为 8 位无符号整数
        image_data_pre = np.uint8(image_data_pre)
        
        # 将处理后的图像添加到 imgs 列表中
        imgs.append(image_data_pre)
        # 确保分割标签中包含的像素点数大于100（这里为啥又问一遍？闻到了屎山的味道）
        assert np.sum(gt_data)>100, 'ground truth should have more than 100 pixels'
        # 将处理后的分割标签添加到gts列表中
        gts.append(gt_data)
        # resize image to 3*1024*1024
        # 创建一个 ResizeLongestSide 对象
        # ResizeLongestSide 的类，用于将图像和坐标进行长边缩放。
        # 具体来说，该类实现了 apply_image 和 apply_coords 两个方法，分别用于处理图像和坐标
        sam_transform = ResizeLongestSide(sam_model.image_encoder.img_size)
        # 将该 ResizeLongestSide 对象应用于 image_data_pre 图像，重新调整大小并返回新的图像 resize_img
        resize_img = sam_transform.apply_image(image_data_pre)
        # 将 numpy 数组 resize_img 转换为 PyTorch 张量，同时将其移动到 GPU
        resize_img_tensor = torch.as_tensor(resize_img.transpose(2, 0, 1)).to(args.device)
        # 对图像进行预处理，例如减去均值、除以标准差等
        input_image = sam_model.preprocess(resize_img_tensor[None,:,:,:]) # (1, 3, 1024, 1024)
        assert input_image.shape == (1, 3, sam_model.image_encoder.img_size, sam_model.image_encoder.img_size), 'input image should be resized to 1024*1024'
        # pre-compute the image embedding
        # 对输入图像进行特征提取，得到图片的 embedding
        with torch.no_grad():
            embedding = sam_model.image_encoder(input_image)
            img_embeddings.append(embedding.cpu().numpy()[0])

# 将所有2D图像以及它们的相关信息（如 ground truth 和 image embedding）保存成一个 npz 文件
if len(imgs)>1:
    imgs = np.stack(imgs, axis=0) # (n, 256, 256, 3) 表示 n 张 256x256 的 RGB 图像
    gts = np.stack(gts, axis=0) # (n, 256, 256) 表示 n 张 256x256 的灰度图像
    img_embeddings = np.stack(img_embeddings, axis=0) # (n, 1, 256, 64, 64) 将每张图像转换为了1个256x64x64的图像embedding
    # 使用np.savez_compressed函数将这三个numpy数组保存成一个npz文件，其中imgs、gts和img_embeddings分别对应三个关键字参数
    np.savez_compressed(join(save_path, args.data_name + '.npz'), imgs=imgs, gts=gts, img_embeddings=img_embeddings)
    # save an example image for sanity check 随机选择一个图像进行可视化检查
    idx = np.random.randint(imgs.shape[0])  # 随机生成一个索引 idx
    # 从 imgs、gts 和 img_embeddings 中提取出该索引对应的图像
    # img_idx、ground truth gt_idx 和 image embedding img_emb_idx
    img_idx = imgs[idx,:,:,:]
    gt_idx = gts[idx,:,:]
    # 代码使用scikit-image库的find_boundaries函数找到gt_idx中每个目标的边缘位置，并将img_idx中边缘位置的像素设为红色
    bd = segmentation.find_boundaries(gt_idx, mode='inner')
    # 将边缘设为红色
    img_idx[bd, :] = [255, 0, 0]
    # 使用io.imsave函数将处理后的img_idx保存成png文件，以便进一步进行可视化检查
    io.imsave(save_path + '.png', img_idx, check_contrast=False)

class NpzDataset(Dataset): 
    def __init__(self, data_root):
        # 读取指定目录下的所有 .npz 文件
        self.data_root = data_root
        self.npz_files = sorted(os.listdir(self.data_root)) 
        self.npz_data = [np.load(join(data_root, f)) for f in self.npz_files]
        # this implementation is ugly but it works (and is also fast for feeding data to GPU) if your server has enough RAM as an alternative, you can also use a list of npy files and load them one by one
        # 这个实现是丑陋的，但它可以工作（并且向 GPU 提供数据的速度也很快）如果你的服务器有足够的 RAM 作为替代方案，你也可以使用 npy 文件列表并一个一个地加载它们
        # 使用了 np.vstack() 函数对这些数组进行垂直方向上的堆叠操作
        # 将它们的 gts 和 img_embeddings 字段整合成两个 numpy 数组: ori_gts 和 img_embeddings
        self.ori_gts = np.vstack([d['gts'] for d in self.npz_data])
        self.img_embeddings = np.vstack([d['img_embeddings'] for d in self.npz_data])
        # 包含了一条有用的调试信息输出语句，输出实际读取数据文件中的 img_embeddings 和 ori_gts 的形状
        print(f"{self.img_embeddings.shape=}, {self.ori_gts.shape=}")
    
    def __len__(self):
        """
        这段代码定义了 __len__ 方法，该方法返回数据集的大小（即所有样本的个数），在该代码中返回的是 ori_gts 数组的第一维大小。
        由于在 NpzDataset 类初始化时，已经将所有 npz 文件中的 gts 字段整合成一个 numpy 数组 ori_gts，因此该方法返回的是所有读取文件中的目标个数（即数据集中的样本数）
        """
        return self.ori_gts.shape[0]
 
    def __getitem__(self, index):
        # 词嵌入向量
        img_embed = self.img_embeddings[index]
        # Ground Truth
        gt2D = self.ori_gts[index]
        # 边界框
        y_indices, x_indices = np.where(gt2D > 0)
        x_min, x_max = np.min(x_indices), np.max(x_indices)
        y_min, y_max = np.min(y_indices), np.max(y_indices)
        # add perturbation to bounding box coordinates
        # 向边界框坐标添加扰动，以实现数据增强
        H, W = gt2D.shape
        x_min = max(0, x_min - np.random.randint(0, 20))
        x_max = min(W, x_max + np.random.randint(0, 20))
        y_min = max(0, y_min - np.random.randint(0, 20))
        y_max = min(H, y_max + np.random.randint(0, 20))
        bboxes = np.array([x_min, y_min, x_max, y_max])
        # convert img embedding, mask, bounding box to torch tensor
        # 返回一个三元组：（一个图像的嵌入向量 img_embed, 对应标注的二维 Ground Truth 图 gt2D, 对应的包含目标的边界框的四个坐标 bboxes）
        return torch.tensor(img_embed).float(), torch.tensor(gt2D[None, :,:]).long(), torch.tensor(bboxes).float()

npz_tr_path = 'data/demo2D_vit_b'
# 使用路径 npz_tr_path 创建了一个新的 NpzDataset 实例 demo_dataset
demo_dataset = NpzDataset(npz_tr_path)
# 训练开始前，代码使用 for 循环从 demo_dataloader 中依次读取一个小批量（batch）的数据，用于测试数据集和数据加载器的正确性。批大小为 8，这意味着每次迭代中将读取 8 个数据样本
demo_dataloader = DataLoader(demo_dataset, batch_size=8, shuffle=True)
for img_embed, gt2D, bboxes in demo_dataloader:
    # img_embed: (B, 256, 64, 64), gt2D: (B, 1, 256, 256), bboxes: (B, 4)
    # 使用 print() 函数打印了从 demo_dataloader 中读取的第一个小批量 img_embed、gt2D 和 bboxes 的形状，以确认它们是否与预期一致
    print(f"{img_embed.shape=}, {gt2D.shape=}, {bboxes.shape=}")
    # 这里程序使用 break 结束了遍历，只输出了第一个小批量的结果
    break

# train data path
npz_tr_path = 'data/demo2D_vit_b'  # 训练数据
work_dir = './work_dir'  # 工作目录路径
task_name = 'demo2D'  # 任务名称
# prepare SAM model
model_type = 'vit_b'  # 模型类型
checkpoint = 'work_dir/SAM/sam_vit_b_01ec64.pth'  # 预训练模型
device = 'cuda:0'  # 设备
model_save_path = join(work_dir, task_name)  # 模型保存地址
os.makedirs(model_save_path, exist_ok=True)
sam_model = sam_model_registry[model_type](checkpoint=checkpoint).to(device)  # 加载模型
sam_model.train()  # 设为训练模式
# Set up the optimizer, hyperparameter tuning will improve performance here
optimizer = torch.optim.Adam(sam_model.mask_decoder.parameters(), lr=1e-5, weight_decay=0)  # 优化器
# 代码定义了一个分割损失函数，其中采用 DiceLoss 和 CrossEntropyLoss 的结合体。
# DiceLoss 是一个测量预测分割与真实分割偏差的指标，CrossEntropyLoss 则是针对多分类问题的损失函数，用于评估预测结果的匹配程度
seg_loss = monai.losses.DiceCELoss(sigmoid=True, squared_pred=True, reduction='mean')  # 损失函数

num_epochs = 100  # 迭代次数
losses = []  # 空列表，用于存放每个 epoch 的损失值
best_loss = 1e10  # 最优损失值
train_dataset = NpzDataset(npz_tr_path)  # 读入训练数据
# 定义数据加载器以便读取和组合数据，同时将样本分成大小为 64 的批次，并打乱顺序
train_dataloader = DataLoader(train_dataset, batch_size=64, shuffle=True)
for epoch in range(num_epochs):
    epoch_loss = 0
    # train
    # step 表示当前处理到了第几个批次
    # image_embedding 是嵌入图像的特征向量
    # gt2D 是训练数据的真实遮罩层标签
    # boxes 是真实的 2D 边界框
    for step, (image_embedding, gt2D, boxes) in enumerate(tqdm(train_dataloader)):
        # do not compute gradients for image encoder and prompt encoder
        # 冻结 图像编码器 和 提示编码器
        with torch.no_grad():
            # convert box to 1024x1024 grid
            # 将边界框坐标从原始坐标系转换为 1024x1024 网格坐标系
            box_np = boxes.numpy()
            # 改变大小
            sam_trans = ResizeLongestSide(sam_model.image_encoder.img_size)
            # 改变提示框大小
            box = sam_trans.apply_boxes(box_np, (gt2D.shape[-2], gt2D.shape[-1]))
            # 转换成 pytorch 张量
            box_torch = torch.as_tensor(box, dtype=torch.float, device=device)
            if len(box_torch.shape) == 2:
                """
                这段代码实现的是获取提示嵌入的过程。
                首先通过 if 语句来判断 box_torch 张量的形状是否为 (B, 4)，
                其中 B 表示批次大小，4 表示边界框的坐标信息（左上角点和右下角点）。
				如果 box_torch 张量的形状是 (B, 4)，则执行 if 语句中的代码进行扩维处理，
				将其转换为形状为 (B, 1, 4) 的张量。
				这么做是为了在后面的计算中保证输入张量的形状一致，从而避免出现维度不匹配的错误
                """
                box_torch = box_torch[:, None, :] # (B, 1, 4)
            # get prompt embeddings 获取提示嵌入
            sparse_embeddings, dense_embeddings = sam_model.prompt_encoder(
                points=None,  # 没有用到点的信息
                boxes=box_torch,  # 使用边界框来提取特征
                masks=None,  # 没有使用遮罩层来进行像素级的聚合
            )
        # predicted masks 前向传播
        mask_predictions, _ = sam_model.mask_decoder(
            image_embeddings=image_embedding.to(device), # (B, 256, 64, 64)
            image_pe=sam_model.prompt_encoder.get_dense_pe(), # (1, 256, 64, 64)
            sparse_prompt_embeddings=sparse_embeddings, # (B, 2, 256)
            dense_prompt_embeddings=dense_embeddings, # (B, 256, 64, 64)
            multimask_output=False,
          )
        
		# 计算损失函数的值
        loss = seg_loss(mask_predictions, gt2D.to(device))
        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        # 记录损失值
        epoch_loss += loss.item()
    
    epoch_loss /= step
    losses.append(epoch_loss)
    print(f'EPOCH: {epoch}, Loss: {epoch_loss}')
    # save the latest model checkpoint
    torch.save(sam_model.state_dict(), join(model_save_path, 'sam_model_latest.pth'))  # 最近一次 checkpoint
    # save the best model
    if epoch_loss < best_loss:
        best_loss = epoch_loss
        torch.save(sam_model.state_dict(), join(model_save_path, 'sam_model_best.pth'))  # 最优 checkpoint

plt.plot(losses)
plt.title('Dice + Cross Entropy Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.show() # comment this line if you are running on a server
plt.savefig(join(model_save_path, 'train_loss.png'))
plt.close()

from skimage import io
# 加载 原始 SAM 模型 到 GPU 上
ori_sam_model = sam_model_registry[model_type](checkpoint=checkpoint).to(device)
# 加载 predictor
ori_sam_predictor = SamPredictor(ori_sam_model)
 
# 读入数据集
ts_img_path = 'data/MedSAMDemo_2D/test/images'
ts_gt_path = 'data/MedSAMDemo_2D/test/labels'
test_names = sorted(os.listdir(ts_img_path))
 
# random select a test case
# 随机读取一张图像
img_idx = np.random.randint(len(test_names))  # 获取索引
image_data = io.imread(join(ts_img_path, test_names[img_idx]))  # 读取
if image_data.shape[-1]>3 and len(image_data.shape)==3:  # 确保图像只有 3 个通道
    image_data = image_data[:,:,:3]
if len(image_data.shape)==2:  # 如果是单通道的灰度图像，转成 3 通道
    image_data = np.repeat(image_data[:,:,None], 3, axis=-1)
# read ground truth (gt should have the same name as the image) and simulate a bounding box
def get_bbox_from_mask(mask):
    '''
    Returns a bounding box from a mask
    从 ground truth 中提取出边界框坐标信息，用于对图像进行裁剪
    '''
    y_indices, x_indices = np.where(mask > 0)
    x_min, x_max = np.min(x_indices), np.max(x_indices)
    y_min, y_max = np.min(y_indices), np.max(y_indices)
    # add perturbation to bounding box coordinates
    H, W = mask.shape
    x_min = max(0, x_min - np.random.randint(0, 20))
    x_max = min(W, x_max + np.random.randint(0, 20))
    y_min = max(0, y_min - np.random.randint(0, 20))
    y_max = min(H, y_max + np.random.randint(0, 20))
 
    return np.array([x_min, y_min, x_max, y_max])
 
# 获得 ground truth
gt_data = io.imread(join(ts_gt_path, test_names[img_idx]))
bbox_raw = get_bbox_from_mask(gt_data)
 
# preprocess: cut-off and max-min normalization 图像预处理
lower_bound, upper_bound = np.percentile(image_data, 0.5), np.percentile(image_data, 99.5)
image_data_pre = np.clip(image_data, lower_bound, upper_bound)
# 亮度范围裁剪
image_data_pre = (image_data_pre - np.min(image_data_pre))/(np.max(image_data_pre)-np.min(image_data_pre))*255.0
image_data_pre[image_data==0] = 0
image_data_pre = np.uint8(image_data_pre)
H, W, _ = image_data_pre.shape
 
# predict the segmentation mask using the original SAM model
# 开跑！
ori_sam_predictor.set_image(image_data_pre)
ori_sam_seg, _, _ = ori_sam_predictor.predict(point_coords=None, box=bbox_raw, multimask_output=False)

# resize image to 3*1024*1024
# 使用 ResizeLongestSide() 函数对原始图像进行大小调整，将其 resize 为 3x1024x1024 的张量
sam_transform = ResizeLongestSide(sam_model.image_encoder.img_size)
resize_img = sam_transform.apply_image(image_data_pre)
# 将调整后的图像张量转换为 PyTorch tensor
resize_img_tensor = torch.as_tensor(resize_img.transpose(2, 0, 1)).to(device)
input_image = sam_model.preprocess(resize_img_tensor[None,:,:,:]) # (1, 3, 1024, 1024)
assert input_image.shape == (1, 3, sam_model.image_encoder.img_size, sam_model.image_encoder.img_size), 'input image should be resized to 1024*1024'
 
with torch.no_grad():
    # pre-compute the image embedding 使用模型的 image_encoder 对象计算图像嵌入向量
    ts_img_embedding = sam_model.image_encoder(input_image)
    # convert box to 1024x1024 grid 将 box 坐标信息调整到 1024x1024 的网络 grid 上
    bbox = sam_trans.apply_boxes(bbox_raw, (H, W))
    print(f'{bbox_raw=} -> {bbox=}')
    box_torch = torch.as_tensor(bbox, dtype=torch.float, device=device)
    if len(box_torch.shape) == 2:
        box_torch = box_torch[:, None, :] # (B, 4) -> (B, 1, 4)
    
    # 使用 prompt_encoder 对象计算稠密和稀疏的嵌入向量（dense and sparse embedding）
    sparse_embeddings, dense_embeddings = sam_model.prompt_encoder(
        points=None,
        boxes=box_torch,
        masks=None,
    )
    medsam_seg_prob, _ = sam_model.mask_decoder(  # 各种图像嵌入向量
        image_embeddings=ts_img_embedding.to(device), # (B, 256, 64, 64)
        image_pe=sam_model.prompt_encoder.get_dense_pe(), # (1, 256, 64, 64)
        sparse_prompt_embeddings=sparse_embeddings, # (B, 2, 256)
        dense_prompt_embeddings=dense_embeddings, # (B, 256, 64, 64)
        multimask_output=False,
        )
    medsam_seg_prob = torch.sigmoid(medsam_seg_prob)  # 压缩到[0, 1]
    # convert soft mask to hard mask
    medsam_seg_prob = medsam_seg_prob.cpu().numpy().squeeze()
    medsam_seg = (medsam_seg_prob > 0.5).astype(np.uint8)
    print(medsam_seg.shape)

ori_sam_dsc = compute_dice_coefficient(gt_data>0, ori_sam_seg>0)
medsam_dsc = compute_dice_coefficient(gt_data>0, medsam_seg>0)
print('Original SAM DSC: {:.4f}'.format(ori_sam_dsc), 'MedSAM DSC: {:.4f}'.format(medsam_dsc))

# visualization functions
# source: https://github.com/facebookresearch/segment-anything/blob/main/notebooks/predictor_example.ipynb
# change color to avoid red and green
def show_mask(mask, ax, random_color=False):
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
    else:
        color = np.array([251/255, 252/255, 30/255, 0.6])
    h, w = mask.shape[-2:]
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    ax.imshow(mask_image)
    
def show_box(box, ax):
    x0, y0 = box[0], box[1]
    w, h = box[2] - box[0], box[3] - box[1]
    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='blue', facecolor=(0,0,0,0), lw=2))    
 
_, axs = plt.subplots(1, 3, figsize=(25, 25))
axs[0].imshow(image_data)
show_mask(gt_data>0, axs[0])
# show_box(box_np[img_id], axs[0])
# axs[0].set_title('Mask with Tuned Model', fontsize=20)
axs[0].axis('off')
 
axs[1].imshow(image_data)
show_mask(ori_sam_seg, axs[1])
show_box(bbox_raw, axs[1])
# add text to image to show dice score
axs[1].text(0.5, 0.5, 'SAM DSC: {:.4f}'.format(ori_sam_dsc), fontsize=30, horizontalalignment='left', verticalalignment='top', color='yellow')
# axs[1].set_title('Mask with Untuned Model', fontsize=20)
axs[1].axis('off')
 
axs[2].imshow(image_data)
show_mask(medsam_seg, axs[2])
show_box(bbox_raw, axs[2])
# add text to image to show dice score
axs[2].text(0.5, 0.5, 'MedSAM DSC: {:.4f}'.format(medsam_dsc), fontsize=30, horizontalalignment='left', verticalalignment='top', color='yellow')
# axs[2].set_title('Ground Truth', fontsize=20)
axs[2].axis('off')
plt.show()  
plt.subplots_adjust(wspace=0.01, hspace=0)
# save plot
# plt.savefig(join(model_save_path, test_npzs[npz_idx].split('.npz')[0] + str(img_id).zfill(3) + '.png'), bbox_inches='tight', dpi=300)
plt.close()