├── README.md └── code ├── segment_anything_lora ├── utils │ ├── __init__.py │ ├── transforms.py │ ├── onnx.py │ └── amg.py ├── modeling │ ├── __init__.py │ ├── common.py │ ├── mask_decoder.py │ ├── sam.py │ ├── transformer.py │ ├── prompt_encoder.py │ └── image_encoder.py ├── __init__.py ├── build_sam.py ├── predictor.py └── automatic_mask_generator.py ├── train.sh ├── utils ├── ramps.py ├── bezier_curve.py └── losses.py ├── networks ├── hierarchical_unet_3d.py └── hierarchical_vnet.py ├── test_singleclass.py ├── test_util_singleclass.py ├── sam_lora_image_encoder.py ├── train_finetuning.py ├── train_retraining.py └── dataloaders └── dataset.py /README.md: -------------------------------------------------------------------------------- 1 | # SFR 2 | 3 | ### Introduction 4 | 5 | This is the Implementation of《Stitching, Fine-tuning, Re-training: A SAM-enabled Framework for Semi-supervised 3D Medical Image Segmentation》 6 | -------------------------------------------------------------------------------- /code/segment_anything_lora/utils/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | -------------------------------------------------------------------------------- /code/train.sh: -------------------------------------------------------------------------------- 1 | cd .. 2 | python3 train_finetuning.py --gpu 0 --exp sam_ft --label_num 16 --patch_size 112 --rdmrotflip 3 | python3 train_retraining.py --gpu 0 --exp 3d_seg --label_num 16 --patch_size 112 --pre_exp sam_ft 4 | # test 5 | python3 test_singleclass.py --gpu 0 --model 3d_seg --iteration 6000 6 | 7 | 8 | 9 | 10 | -------------------------------------------------------------------------------- /code/segment_anything_lora/modeling/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | from .sam import Sam 8 | from .image_encoder import ImageEncoderViT 9 | from .mask_decoder import MaskDecoder 10 | from .prompt_encoder import PromptEncoder 11 | from .transformer import TwoWayTransformer 12 | -------------------------------------------------------------------------------- /code/segment_anything_lora/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | from .build_sam import ( 8 | build_sam, 9 | build_sam_vit_h, 10 | build_sam_vit_l, 11 | build_sam_vit_b, 12 | sam_model_registry, 13 | ) 14 | from .predictor import SamPredictor 15 | from .automatic_mask_generator import SamAutomaticMaskGenerator 16 | -------------------------------------------------------------------------------- /code/utils/ramps.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018, Curious AI Ltd. All rights reserved. 2 | # 3 | # This work is licensed under the Creative Commons Attribution-NonCommercial 4 | # 4.0 International License. To view a copy of this license, visit 5 | # http://creativecommons.org/licenses/by-nc/4.0/ or send a letter to 6 | # Creative Commons, PO Box 1866, Mountain View, CA 94042, USA. 7 | 8 | """Functions for ramping hyperparameters up or down 9 | 10 | Each function takes the current training step or epoch, and the 11 | ramp length in the same format, and returns a multiplier between 12 | 0 and 1. 13 | """ 14 | 15 | 16 | import numpy as np 17 | 18 | 19 | def sigmoid_rampup(current, rampup_length): 20 | """Exponential rampup from https://arxiv.org/abs/1610.02242""" 21 | if rampup_length == 0: 22 | return 1.0 23 | else: 24 | current = np.clip(current, 0.0, rampup_length) 25 | phase = 1.0 - current / rampup_length 26 | return float(np.exp(-5.0 * phase * phase)) 27 | 28 | 29 | def linear_rampup(current, rampup_length): 30 | """Linear rampup""" 31 | assert current >= 0 and rampup_length >= 0 32 | if current >= rampup_length: 33 | return 1.0 34 | else: 35 | return current / rampup_length 36 | 37 | 38 | def cosine_rampdown(current, rampdown_length): 39 | """Cosine rampdown from https://arxiv.org/abs/1608.03983""" 40 | assert 0 <= current <= rampdown_length 41 | return float(.5 * (np.cos(np.pi * current / rampdown_length) + 1)) 42 | -------------------------------------------------------------------------------- /code/segment_anything_lora/modeling/common.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import torch 8 | import torch.nn as nn 9 | 10 | from typing import Type 11 | 12 | 13 | class MLPBlock(nn.Module): 14 | def __init__( 15 | self, 16 | embedding_dim: int, 17 | mlp_dim: int, 18 | act: Type[nn.Module] = nn.GELU, 19 | ) -> None: 20 | super().__init__() 21 | self.lin1 = nn.Linear(embedding_dim, mlp_dim) 22 | self.lin2 = nn.Linear(mlp_dim, embedding_dim) 23 | self.act = act() 24 | 25 | def forward(self, x: torch.Tensor) -> torch.Tensor: 26 | return self.lin2(self.act(self.lin1(x))) 27 | 28 | 29 | # From https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/batch_norm.py # noqa 30 | # Itself from https://github.com/facebookresearch/ConvNeXt/blob/d1fa8f6fef0a165b27399986cc2bdacc92777e40/models/convnext.py#L119 # noqa 31 | class LayerNorm2d(nn.Module): 32 | def __init__(self, num_channels: int, eps: float = 1e-6) -> None: 33 | super().__init__() 34 | self.weight = nn.Parameter(torch.ones(num_channels)) 35 | self.bias = nn.Parameter(torch.zeros(num_channels)) 36 | self.eps = eps 37 | 38 | def forward(self, x: torch.Tensor) -> torch.Tensor: 39 | u = x.mean(1, keepdim=True) 40 | s = (x - u).pow(2).mean(1, keepdim=True) 41 | x = (x - u) / torch.sqrt(s + self.eps) 42 | x = self.weight[:, None, None] * x + self.bias[:, None, None] 43 | return x 44 | -------------------------------------------------------------------------------- /code/utils/bezier_curve.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import random 3 | # import matplotlib.pyplot as plt 4 | try: 5 | from scipy.special import comb 6 | except: 7 | from scipy.misc import comb 8 | 9 | 10 | def bernstein_poly(i, n, t): 11 | """ 12 | The Bernstein polynomial of n, i as a function of t 13 | """ 14 | return comb(n, i) * ( t**(n-i) ) * (1 - t)**i 15 | 16 | 17 | def bezier_curve(points, nTimes=1000): 18 | """ 19 | Given a set of control points, return the 20 | bezier curve defined by the control points. 21 | Control points should be a list of lists, or list of tuples 22 | such as [ [1,1], 23 | [2,3], 24 | [4,5], ..[Xn, Yn] ] 25 | nTimes is the number of time steps, defaults to 1000 26 | See http://processingjs.nihongoresources.com/bezierinfo/ 27 | """ 28 | 29 | nPoints = len(points) 30 | xPoints = np.array([p[0] for p in points]) 31 | yPoints = np.array([p[1] for p in points]) 32 | 33 | t = np.linspace(0.0, 1.0, nTimes) 34 | 35 | polynomial_array = np.array([bernstein_poly(i, nPoints-1, t) for i in range(0, nPoints)]) 36 | 37 | xvals = np.dot(xPoints, polynomial_array) 38 | yvals = np.dot(yPoints, polynomial_array) 39 | 40 | return xvals, yvals 41 | 42 | 43 | def nonlinear_transformation_r1(x, prob=0.5): 44 | if random.random() >= prob: 45 | return x 46 | points = [[0, 0], [random.random(), random.random()], [random.random(), random.random()], [1, 1]] 47 | xvals, yvals = bezier_curve(points, nTimes=100000) 48 | if random.random() < 0.5: 49 | # Half change to get flip 50 | xvals = np.sort(xvals) 51 | else: 52 | xvals, yvals = np.sort(xvals), np.sort(yvals) 53 | nonlinear_x = np.interp(x, xvals, yvals) 54 | return nonlinear_x 55 | 56 | 57 | def nonlinear_transformation_r2(x, prob=0.5): 58 | if random.random() >= prob: 59 | return x 60 | points = [[0, 0], [random.random(), random.random()], [random.random(), random.random()], [1, 1]] 61 | xvals, yvals = bezier_curve(points, nTimes=100000) 62 | xvals, yvals = np.sort(xvals), np.sort(yvals) 63 | nonlinear_x = np.interp(x, xvals, yvals) 64 | return nonlinear_x 65 | 66 | 67 | def nonlinear_transformation_r2_t(x, prob=0.5): 68 | points = [[0, 0], [random.random(), random.random()], [random.random(), random.random()], [1, 1]] 69 | xvals, yvals = bezier_curve(points, nTimes=100000) 70 | xvals, yvals = np.sort(xvals), np.sort(yvals) 71 | nonlinear_x = np.interp(x, xvals, yvals) 72 | return nonlinear_x 73 | 74 | 75 | def nonlinear_transformation_r3(x, t1, t2, t3, t4): 76 | points = [[0, 0], [t1, t2], [t3, t4], [1, 1]] 77 | xvals, yvals = bezier_curve(points, nTimes=100000) 78 | if random.random() < 0.5: 79 | # Half change to get flip 80 | xvals = np.sort(xvals) 81 | else: 82 | xvals, yvals = np.sort(xvals), np.sort(yvals) 83 | nonlinear_x = np.interp(x, xvals, yvals) 84 | return nonlinear_x 85 | 86 | 87 | def nonlinear_transformation_r4(x): 88 | if random.random() >= 0.5: 89 | return x 90 | if random.random() < 0.5: 91 | points = [[-1, -1], [-0.5, 0.5], [0.5, -0.5], [1, 1]] 92 | else: 93 | points = [[-1, -1], [-0.75, 0.75], [0.75, -0.75], [1, 1]] 94 | xvals, yvals = bezier_curve(points, nTimes=100000) 95 | xvals = np.sort(xvals) 96 | yvals = np.sort(yvals) 97 | 98 | nonlinear_x = np.interp(x, xvals, yvals) 99 | 100 | return nonlinear_x -------------------------------------------------------------------------------- /code/segment_anything_lora/utils/transforms.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import numpy as np 8 | import torch 9 | from torch.nn import functional as F 10 | from torchvision.transforms.functional import resize, to_pil_image # type: ignore 11 | 12 | from copy import deepcopy 13 | from typing import Tuple 14 | 15 | 16 | class ResizeLongestSide: 17 | """ 18 | Resizes images to longest side 'target_length', as well as provides 19 | methods for resizing coordinates and boxes. Provides methods for 20 | transforming both numpy array and batched torch tensors. 21 | """ 22 | 23 | def __init__(self, target_length: int) -> None: 24 | self.target_length = target_length 25 | 26 | def apply_image(self, image: np.ndarray) -> np.ndarray: 27 | """ 28 | Expects a numpy array with shape HxWxC in uint8 format. 29 | """ 30 | target_size = self.get_preprocess_shape(image.shape[0], image.shape[1], self.target_length) 31 | return np.array(resize(to_pil_image(image), target_size)) 32 | 33 | def apply_coords(self, coords: np.ndarray, original_size: Tuple[int, ...]) -> np.ndarray: 34 | """ 35 | Expects a numpy array of length 2 in the final dimension. Requires the 36 | original image size in (H, W) format. 37 | """ 38 | old_h, old_w = original_size 39 | new_h, new_w = self.get_preprocess_shape( 40 | original_size[0], original_size[1], self.target_length 41 | ) 42 | coords = deepcopy(coords).astype(float) 43 | coords[..., 0] = coords[..., 0] * (new_w / old_w) 44 | coords[..., 1] = coords[..., 1] * (new_h / old_h) 45 | return coords 46 | 47 | def apply_boxes(self, boxes: np.ndarray, original_size: Tuple[int, ...]) -> np.ndarray: 48 | """ 49 | Expects a numpy array shape Bx4. Requires the original image size 50 | in (H, W) format. 51 | """ 52 | boxes = self.apply_coords(boxes.reshape(-1, 2, 2), original_size) 53 | return boxes.reshape(-1, 4) 54 | 55 | def apply_image_torch(self, image: torch.Tensor) -> torch.Tensor: 56 | """ 57 | Expects batched images with shape BxCxHxW and float format. This 58 | transformation may not exactly match apply_image. apply_image is 59 | the transformation expected by the model. 60 | """ 61 | # Expects an image in BCHW format. May not exactly match apply_image. 62 | target_size = self.get_preprocess_shape(image.shape[0], image.shape[1], self.target_length) 63 | return F.interpolate( 64 | image, target_size, mode="bilinear", align_corners=False, antialias=True 65 | ) 66 | 67 | def apply_coords_torch( 68 | self, coords: torch.Tensor, original_size: Tuple[int, ...] 69 | ) -> torch.Tensor: 70 | """ 71 | Expects a torch tensor with length 2 in the last dimension. Requires the 72 | original image size in (H, W) format. 73 | """ 74 | old_h, old_w = original_size 75 | new_h, new_w = self.get_preprocess_shape( 76 | original_size[0], original_size[1], self.target_length 77 | ) 78 | coords = deepcopy(coords).to(torch.float) 79 | coords[..., 0] = coords[..., 0] * (new_w / old_w) 80 | coords[..., 1] = coords[..., 1] * (new_h / old_h) 81 | return coords 82 | 83 | def apply_boxes_torch( 84 | self, boxes: torch.Tensor, original_size: Tuple[int, ...] 85 | ) -> torch.Tensor: 86 | """ 87 | Expects a torch tensor with shape Bx4. Requires the original image 88 | size in (H, W) format. 89 | """ 90 | boxes = self.apply_coords_torch(boxes.reshape(-1, 2, 2), original_size) 91 | return boxes.reshape(-1, 4) 92 | 93 | @staticmethod 94 | def get_preprocess_shape(oldh: int, oldw: int, long_side_length: int) -> Tuple[int, int]: 95 | """ 96 | Compute the output size given input size and target long side length. 97 | """ 98 | scale = long_side_length * 1.0 / max(oldh, oldw) 99 | newh, neww = oldh * scale, oldw * scale 100 | neww = int(neww + 0.5) 101 | newh = int(newh + 0.5) 102 | return (newh, neww) 103 | -------------------------------------------------------------------------------- /code/networks/hierarchical_unet_3d.py: -------------------------------------------------------------------------------- 1 | """ 2 | =unet_3D_dv_semi.py，是unet_3D.py的深监督版本 3 | https://github.com/HiLab-git/SSL4MIS/blob/master/code/networks/unet_3D_dv_semi.py 4 | This file is adapted from https://github.com/ozan-oktay/Attention-Gated-Networks 5 | """ 6 | 7 | import math 8 | import torch 9 | import torch.nn as nn 10 | from networks.utils import UnetConv3, UnetUp3, UnetUp3_CT, UnetDsv3 11 | import torch.nn.functional as F 12 | from networks.networks_other import init_weights 13 | 14 | 15 | class UNet_3D(nn.Module): 16 | 17 | def __init__(self, feature_scale=4, n_classes=21, is_deconv=True, in_channels=3, is_batchnorm=True): 18 | super(UNet_3D, self).__init__() 19 | self.is_deconv = is_deconv 20 | self.in_channels = in_channels 21 | self.is_batchnorm = is_batchnorm 22 | self.feature_scale = feature_scale 23 | 24 | filters = [64, 128, 256, 512, 1024] 25 | filters = [int(x / self.feature_scale) for x in filters] 26 | 27 | # downsampling 28 | self.conv1 = UnetConv3(self.in_channels, filters[0], self.is_batchnorm, kernel_size=( 29 | 3, 3, 3), padding_size=(1, 1, 1)) 30 | self.maxpool1 = nn.MaxPool3d(kernel_size=(2, 2, 2)) 31 | 32 | self.conv2 = UnetConv3(filters[0], filters[1], self.is_batchnorm, kernel_size=( 33 | 3, 3, 3), padding_size=(1, 1, 1)) 34 | self.maxpool2 = nn.MaxPool3d(kernel_size=(2, 2, 2)) 35 | 36 | self.conv3 = UnetConv3(filters[1], filters[2], self.is_batchnorm, kernel_size=( 37 | 3, 3, 3), padding_size=(1, 1, 1)) 38 | self.maxpool3 = nn.MaxPool3d(kernel_size=(2, 2, 2)) 39 | 40 | self.conv4 = UnetConv3(filters[2], filters[3], self.is_batchnorm, kernel_size=( 41 | 3, 3, 3), padding_size=(1, 1, 1)) 42 | self.maxpool4 = nn.MaxPool3d(kernel_size=(2, 2, 2)) 43 | 44 | self.center = UnetConv3(filters[3], filters[4], self.is_batchnorm, kernel_size=( 45 | 3, 3, 3), padding_size=(1, 1, 1)) 46 | 47 | # upsampling 48 | self.up_concat4 = UnetUp3_CT(filters[4], filters[3], is_batchnorm) 49 | self.up_concat3 = UnetUp3_CT(filters[3], filters[2], is_batchnorm) 50 | self.up_concat2 = UnetUp3_CT(filters[2], filters[1], is_batchnorm) 51 | self.up_concat1 = UnetUp3_CT(filters[1], filters[0], is_batchnorm) 52 | 53 | # deep supervision 54 | self.dsv4 = UnetDsv3(in_size=filters[3], out_size=n_classes, scale_factor=8) 55 | self.dsv3 = UnetDsv3(in_size=filters[2], out_size=n_classes, scale_factor=4) 56 | self.dsv2 = UnetDsv3(in_size=filters[1], out_size=n_classes, scale_factor=2) 57 | self.dsv1 = nn.Conv3d(in_channels=filters[0], out_channels=n_classes, kernel_size=1) 58 | 59 | self.dropout1 = nn.Dropout3d(p=0.5) 60 | self.dropout2 = nn.Dropout3d(p=0.3) 61 | self.dropout3 = nn.Dropout3d(p=0.2) 62 | self.dropout4 = nn.Dropout3d(p=0.1) 63 | 64 | # initialise weights 65 | for m in self.modules(): 66 | if isinstance(m, nn.Conv3d): 67 | init_weights(m, init_type='kaiming') 68 | elif isinstance(m, nn.BatchNorm3d): 69 | init_weights(m, init_type='kaiming') 70 | 71 | def forward(self, inputs): 72 | conv1 = self.conv1(inputs) 73 | maxpool1 = self.maxpool1(conv1) 74 | 75 | conv2 = self.conv2(maxpool1) 76 | maxpool2 = self.maxpool2(conv2) 77 | 78 | conv3 = self.conv3(maxpool2) 79 | maxpool3 = self.maxpool3(conv3) 80 | 81 | conv4 = self.conv4(maxpool3) 82 | maxpool4 = self.maxpool4(conv4) 83 | 84 | center = self.center(maxpool4) 85 | 86 | up4 = self.up_concat4(conv4, center) 87 | up4 = self.dropout1(up4) 88 | 89 | up3 = self.up_concat3(conv3, up4) 90 | up3 = self.dropout2(up3) 91 | 92 | up2 = self.up_concat2(conv2, up3) 93 | up2 = self.dropout3(up2) 94 | 95 | up1 = self.up_concat1(conv1, up2) 96 | up1 = self.dropout4(up1) 97 | 98 | # Deep Supervision 99 | dsv4 = self.dsv4(up4) 100 | dsv3 = self.dsv3(up3) 101 | dsv2 = self.dsv2(up2) 102 | dsv1 = self.dsv1(up1) 103 | 104 | return dsv1, dsv2, dsv3, dsv4 105 | 106 | @staticmethod 107 | def apply_argmax_softmax(pred): 108 | log_p = F.softmax(pred, dim=1) 109 | 110 | return log_p -------------------------------------------------------------------------------- /code/test_singleclass.py: -------------------------------------------------------------------------------- 1 | import os 2 | import argparse 3 | import torch 4 | import numpy as np 5 | from networks.hierarchical_vnet import VNet 6 | from networks.hierarchical_unet_3d import UNet_3D 7 | from test_util_singleclass_hv import test_all_case 8 | import pdb 9 | 10 | parser = argparse.ArgumentParser() 11 | parser.add_argument('--root_path', type=str, default='../data/LA/2018LA_Seg_Training Set/', help='Name of Experiment') 12 | parser.add_argument('--model', type=str, default='sparse_mt_294', help='model_name') 13 | parser.add_argument('--dataset', type=str, default='la', help='dataset to use') 14 | parser.add_argument('--data_version', type=str, default='v2', help='dataset version to use') 15 | parser.add_argument('--set_version', type=str, default='0', help='dataset version to use') 16 | parser.add_argument('--semantic_class', type=str, default='kidney', choices=['kidney', 'tumor']) 17 | parser.add_argument('--list_num', type=str, default='', help='data list to use') 18 | parser.add_argument('--gpu', type=str, default='0', help='GPU to use') 19 | parser.add_argument('--iteration', type=int, default=6000, help='GPU to use') 20 | parser.add_argument('--patch_size', type=int, default=112, help='patch size') 21 | parser.add_argument('--model_type', type=str, default='vnet', help='model_type') 22 | args = parser.parse_args() 23 | 24 | root = "../" 25 | 26 | os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu 27 | snapshot_path = root + "model_" + args.dataset + "/" + args.model + "/" 28 | test_save_path = root + "model_" + args.dataset + "/prediction/" + args.model + "_post/" 29 | if not os.path.exists(test_save_path): 30 | os.makedirs(test_save_path) 31 | 32 | num_classes = 2 33 | 34 | if args.dataset == 'la': 35 | with open(args.root_path + '/../test' + args.list_num + '.list', 'r') as f: 36 | image_list = f.readlines() 37 | image_list = [args.root_path + item.replace('\n', '') + "/mri_norm2.h5" for item in image_list] 38 | elif args.dataset == 'btcv': 39 | num_classes = 14 40 | with open(args.root_path + '/../test' + args.list_num + '.list', 'r') as f: 41 | image_list = f.readlines() 42 | image_list = [args.root_path + '/' + item.replace('\n', '') + ".h5" for item in image_list] 43 | elif args.dataset == 'mact': 44 | num_classes = 9 45 | with open(args.root_path + '/../test' + args.list_num + '.list', 'r') as f: 46 | image_list = f.readlines() 47 | image_list = [args.root_path + '/' + item.replace('\n', '') + ".h5" for item in image_list] 48 | elif args.dataset == 'brats': 49 | with open(args.root_path + '/../test_follow.list', 'r') as f: 50 | image_list = f.readlines() 51 | image_list = [args.root_path + '/' + item.replace('\n', '') + ".h5" for item in image_list] 52 | 53 | 54 | def test_calculate_metric(epoch_num): 55 | if args.model_type == "vnet": 56 | net = VNet(n_channels=1, n_classes=num_classes, normalization='batchnorm', has_dropout=False, pyramid_has_dropout=False).cuda() 57 | elif args.model_type == "unet_3d": 58 | net = UNet_3D(in_channels=1, n_classes=num_classes).cuda() 59 | save_mode_path = os.path.join(snapshot_path, 'iter_' + str(epoch_num) + '.pth') 60 | net.load_state_dict(torch.load(save_mode_path)) 61 | print("init weight from {}".format(save_mode_path)) 62 | net.eval() 63 | 64 | if args.dataset == 'la': 65 | if args.patch_size == 112: 66 | ps = (112, 112, 80) 67 | elif args.patch_size == 128: 68 | ps = (128, 128, 64) 69 | avg_metric = test_all_case(net, args.dataset, args.semantic_class, image_list, num_classes=num_classes, patch_size=ps, 70 | save_result=True, stride_xy=18, stride_z=4, test_save_path=test_save_path) 71 | 72 | elif args.dataset == 'btcv' or args.dataset == 'mact': 73 | patch_size = args.patch_size 74 | avg_metric = test_all_case(net, args.dataset, args.semantic_class, image_list, num_classes=num_classes, patch_size=(patch_size, patch_size, patch_size), 75 | save_result=True, stride_xy=12, stride_z=12, test_save_path=test_save_path) 76 | 77 | elif args.dataset == 'brats': 78 | patch_size = args.patch_size 79 | avg_metric = test_all_case(net, args.dataset, args.semantic_class, image_list, num_classes=num_classes, patch_size=(patch_size, patch_size, patch_size), 80 | save_result=True, stride_xy=64, stride_z=64, test_save_path=test_save_path) 81 | 82 | return avg_metric 83 | 84 | 85 | if __name__ == '__main__': 86 | metric, std = test_calculate_metric(args.iteration) 87 | print(metric) 88 | with open(root + "model_" + args.dataset + "/prediction_v2.txt", "a") as f: 89 | f.write(args.model + " - " + str(args.iteration) + ": " + ", ".join(str(i) for i in metric) + "\n") 90 | f.write(args.model + " - " + str(args.iteration) + ": " + ", ".join(str(i) for i in std) + "\n") 91 | -------------------------------------------------------------------------------- /code/test_util_singleclass.py: -------------------------------------------------------------------------------- 1 | import pdb 2 | import h5py 3 | import math 4 | import nibabel as nib 5 | import numpy as np 6 | from medpy import metric 7 | import torch 8 | import torch.nn.functional as F 9 | from tqdm import tqdm 10 | from scipy import ndimage 11 | 12 | 13 | def test_all_case(net, dataset, semantic_class, image_list, num_classes, patch_size=(112, 112, 80), stride_xy=18, stride_z=4, save_result=True, test_save_path=None, preproc_fn=None): 14 | total_metric = 0.0 15 | total_array = np.zeros((len(image_list), 4)) 16 | i = 0 17 | # total_metric = [] 18 | for image_path in tqdm(image_list): 19 | h5f = h5py.File(image_path, 'r') 20 | image = h5f['image'][:] 21 | label = h5f['label'][:] 22 | 23 | if dataset == "la": 24 | id = image_path.split('/')[-2] 25 | image = (image - np.mean(image)) / np.std(image) 26 | elif dataset == "pancreas" or dataset == 'btcv' or dataset == 'mact': 27 | id = image_path.split('.')[-2].split('/')[-1] 28 | image = (image - np.mean(image)) / np.std(image) 29 | elif dataset == "lits" or dataset == "kits" or dataset == "promise" or dataset == "acdc" or dataset == "brats": 30 | id = image_path.split('.')[-2].split('/')[-1] 31 | image = image.swapaxes(0, 2) 32 | label = label.swapaxes(0, 2) 33 | image = (image - np.mean(image)) / np.std(image) 34 | if semantic_class == 'tumor': 35 | label[label != 2] = 0 36 | label[label == 2] = 1 37 | else: 38 | label[label > 0] = 1 39 | 40 | if preproc_fn is not None: 41 | image = preproc_fn(image) 42 | prediction, score_map = test_single_case(net, image, stride_xy, stride_z, patch_size, num_classes=num_classes) 43 | 44 | if np.sum(prediction)==0: 45 | single_metric = (0,0,0,0) 46 | else: 47 | single_metric = calculate_metric_percase(prediction, label[:]) # (175, 132, 88), (175, 132, 88) 48 | print(id, single_metric) 49 | total_metric += np.asarray(single_metric) 50 | total_array[i] = single_metric 51 | i += 1 52 | 53 | if save_result: 54 | nib.save(nib.Nifti1Image(prediction.astype(np.float32), np.eye(4)), test_save_path + id + "_pred.nii.gz") 55 | nib.save(nib.Nifti1Image(score_map[1].astype(np.float32), np.eye(4)), test_save_path + id + "_score.nii.gz") 56 | nib.save(nib.Nifti1Image(image[:].astype(np.float32), np.eye(4)), test_save_path + id + "_img.nii.gz") 57 | nib.save(nib.Nifti1Image(label[:].astype(np.float32), np.eye(4)), test_save_path + id + "_gt.nii.gz") 58 | 59 | avg_metric = total_metric / len(image_list) 60 | std = np.std(total_array, axis=0, ddof=1) 61 | print('average metric is {}'.format(avg_metric)) 62 | print('average std is {}'.format(std)) 63 | 64 | return avg_metric, std 65 | 66 | 67 | def test_single_case(net, image, stride_xy, stride_z, patch_size, num_classes=1): 68 | w, h, d = image.shape 69 | 70 | # if the size of image is less than patch_size, then padding it 71 | add_pad = False 72 | if w < patch_size[0]: 73 | w_pad = patch_size[0]-w 74 | add_pad = True 75 | else: 76 | w_pad = 0 77 | if h < patch_size[1]: 78 | h_pad = patch_size[1]-h 79 | add_pad = True 80 | else: 81 | h_pad = 0 82 | if d < patch_size[2]: 83 | d_pad = patch_size[2]-d 84 | add_pad = True 85 | else: 86 | d_pad = 0 87 | wl_pad, wr_pad = w_pad//2,w_pad-w_pad//2 88 | hl_pad, hr_pad = h_pad//2,h_pad-h_pad//2 89 | dl_pad, dr_pad = d_pad//2,d_pad-d_pad//2 90 | if add_pad: 91 | image = np.pad(image, [(wl_pad,wr_pad),(hl_pad,hr_pad), (dl_pad, dr_pad)], mode='constant', constant_values=0) 92 | ww,hh,dd = image.shape 93 | 94 | sx = math.ceil((ww - patch_size[0]) / stride_xy) + 1 95 | sy = math.ceil((hh - patch_size[1]) / stride_xy) + 1 96 | sz = math.ceil((dd - patch_size[2]) / stride_z) + 1 97 | # print("{}, {}, {}".format(sx, sy, sz)) 98 | score_map = np.zeros((num_classes, ) + image.shape).astype(np.float32) 99 | cnt = np.zeros(image.shape).astype(np.float32) 100 | 101 | for x in range(0, sx): 102 | xs = min(stride_xy*x, ww-patch_size[0]) 103 | for y in range(0, sy): 104 | ys = min(stride_xy * y,hh-patch_size[1]) 105 | for z in range(0, sz): 106 | zs = min(stride_z * z, dd-patch_size[2]) 107 | test_patch = image[xs:xs+patch_size[0], ys:ys+patch_size[1], zs:zs+patch_size[2]] 108 | test_patch = np.expand_dims(np.expand_dims(test_patch,axis=0),axis=0).astype(np.float32) 109 | test_patch = torch.from_numpy(test_patch).cuda() 110 | y1, _, _, _ = net(test_patch) # [1, 2, 112, 112, 80] 111 | y = F.softmax(y1, dim=1) 112 | y = y.cpu().data.numpy() 113 | y = y[0,:,:,:,:] 114 | score_map[:, xs:xs+patch_size[0], ys:ys+patch_size[1], zs:zs+patch_size[2]] \ 115 | = score_map[:, xs:xs+patch_size[0], ys:ys+patch_size[1], zs:zs+patch_size[2]] + y 116 | cnt[xs:xs+patch_size[0], ys:ys+patch_size[1], zs:zs+patch_size[2]] \ 117 | = cnt[xs:xs+patch_size[0], ys:ys+patch_size[1], zs:zs+patch_size[2]] + 1 118 | score_map = score_map/np.expand_dims(cnt,axis=0) 119 | label_map = np.argmax(score_map, axis = 0) 120 | if add_pad: 121 | label_map = label_map[wl_pad:wl_pad+w,hl_pad:hl_pad+h,dl_pad:dl_pad+d] 122 | score_map = score_map[:,wl_pad:wl_pad+w,hl_pad:hl_pad+h,dl_pad:dl_pad+d] 123 | return label_map, score_map 124 | 125 | 126 | def calculate_metric_percase(pred, gt): 127 | dice = metric.binary.dc(pred, gt) 128 | jc = metric.binary.jc(pred, gt) 129 | hd = metric.binary.hd95(pred, gt) 130 | asd = metric.binary.asd(pred, gt) 131 | 132 | return dice, jc, hd, asd 133 | -------------------------------------------------------------------------------- /code/segment_anything_lora/utils/onnx.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import torch 8 | import torch.nn as nn 9 | from torch.nn import functional as F 10 | 11 | from typing import Tuple 12 | 13 | from ..modeling import Sam 14 | from .amg import calculate_stability_score 15 | 16 | 17 | class SamOnnxModel(nn.Module): 18 | """ 19 | This model should not be called directly, but is used in ONNX export. 20 | It combines the prompt encoder, mask decoder, and mask postprocessing of Sam, 21 | with some functions modified to enable model tracing. Also supports extra 22 | options controlling what information. See the ONNX export script for details. 23 | """ 24 | 25 | def __init__( 26 | self, 27 | model: Sam, 28 | return_single_mask: bool, 29 | use_stability_score: bool = False, 30 | return_extra_metrics: bool = False, 31 | ) -> None: 32 | super().__init__() 33 | self.mask_decoder = model.mask_decoder 34 | self.model = model 35 | self.img_size = model.image_encoder.img_size 36 | self.return_single_mask = return_single_mask 37 | self.use_stability_score = use_stability_score 38 | self.stability_score_offset = 1.0 39 | self.return_extra_metrics = return_extra_metrics 40 | 41 | @staticmethod 42 | def resize_longest_image_size( 43 | input_image_size: torch.Tensor, longest_side: int 44 | ) -> torch.Tensor: 45 | input_image_size = input_image_size.to(torch.float32) 46 | scale = longest_side / torch.max(input_image_size) 47 | transformed_size = scale * input_image_size 48 | transformed_size = torch.floor(transformed_size + 0.5).to(torch.int64) 49 | return transformed_size 50 | 51 | def _embed_points(self, point_coords: torch.Tensor, point_labels: torch.Tensor) -> torch.Tensor: 52 | point_coords = point_coords + 0.5 53 | point_coords = point_coords / self.img_size 54 | point_embedding = self.model.prompt_encoder.pe_layer._pe_encoding(point_coords) 55 | point_labels = point_labels.unsqueeze(-1).expand_as(point_embedding) 56 | 57 | point_embedding = point_embedding * (point_labels != -1) 58 | point_embedding = point_embedding + self.model.prompt_encoder.not_a_point_embed.weight * ( 59 | point_labels == -1 60 | ) 61 | 62 | for i in range(self.model.prompt_encoder.num_point_embeddings): 63 | point_embedding = point_embedding + self.model.prompt_encoder.point_embeddings[ 64 | i 65 | ].weight * (point_labels == i) 66 | 67 | return point_embedding 68 | 69 | def _embed_masks(self, input_mask: torch.Tensor, has_mask_input: torch.Tensor) -> torch.Tensor: 70 | mask_embedding = has_mask_input * self.model.prompt_encoder.mask_downscaling(input_mask) 71 | mask_embedding = mask_embedding + ( 72 | 1 - has_mask_input 73 | ) * self.model.prompt_encoder.no_mask_embed.weight.reshape(1, -1, 1, 1) 74 | return mask_embedding 75 | 76 | def mask_postprocessing(self, masks: torch.Tensor, orig_im_size: torch.Tensor) -> torch.Tensor: 77 | masks = F.interpolate( 78 | masks, 79 | size=(self.img_size, self.img_size), 80 | mode="bilinear", 81 | align_corners=False, 82 | ) 83 | 84 | prepadded_size = self.resize_longest_image_size(orig_im_size, self.img_size) 85 | masks = masks[..., : int(prepadded_size[0]), : int(prepadded_size[1])] 86 | 87 | orig_im_size = orig_im_size.to(torch.int64) 88 | h, w = orig_im_size[0], orig_im_size[1] 89 | masks = F.interpolate(masks, size=(h, w), mode="bilinear", align_corners=False) 90 | return masks 91 | 92 | def select_masks( 93 | self, masks: torch.Tensor, iou_preds: torch.Tensor, num_points: int 94 | ) -> Tuple[torch.Tensor, torch.Tensor]: 95 | # Determine if we should return the multiclick mask or not from the number of points. 96 | # The reweighting is used to avoid control flow. 97 | score_reweight = torch.tensor( 98 | [[1000] + [0] * (self.model.mask_decoder.num_mask_tokens - 1)] 99 | ).to(iou_preds.device) 100 | score = iou_preds + (num_points - 2.5) * score_reweight 101 | best_idx = torch.argmax(score, dim=1) 102 | masks = masks[torch.arange(masks.shape[0]), best_idx, :, :].unsqueeze(1) 103 | iou_preds = iou_preds[torch.arange(masks.shape[0]), best_idx].unsqueeze(1) 104 | 105 | return masks, iou_preds 106 | 107 | @torch.no_grad() 108 | def forward( 109 | self, 110 | image_embeddings: torch.Tensor, 111 | point_coords: torch.Tensor, 112 | point_labels: torch.Tensor, 113 | mask_input: torch.Tensor, 114 | has_mask_input: torch.Tensor, 115 | orig_im_size: torch.Tensor, 116 | ): 117 | sparse_embedding = self._embed_points(point_coords, point_labels) 118 | dense_embedding = self._embed_masks(mask_input, has_mask_input) 119 | 120 | masks, scores = self.model.mask_decoder.predict_masks( 121 | image_embeddings=image_embeddings, 122 | image_pe=self.model.prompt_encoder.get_dense_pe(), 123 | sparse_prompt_embeddings=sparse_embedding, 124 | dense_prompt_embeddings=dense_embedding, 125 | ) 126 | 127 | if self.use_stability_score: 128 | scores = calculate_stability_score( 129 | masks, self.model.mask_threshold, self.stability_score_offset 130 | ) 131 | 132 | if self.return_single_mask: 133 | masks, scores = self.select_masks(masks, scores, point_coords.shape[1]) 134 | 135 | upscaled_masks = self.mask_postprocessing(masks, orig_im_size) 136 | 137 | if self.return_extra_metrics: 138 | stability_scores = calculate_stability_score( 139 | upscaled_masks, self.model.mask_threshold, self.stability_score_offset 140 | ) 141 | areas = (upscaled_masks > self.model.mask_threshold).sum(-1).sum(-1) 142 | return upscaled_masks, scores, stability_scores, areas, masks 143 | 144 | return upscaled_masks, scores, masks 145 | -------------------------------------------------------------------------------- /code/segment_anything_lora/build_sam.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import torch 8 | from torch.nn import functional as F 9 | from functools import partial 10 | from .modeling import ImageEncoderViT, MaskDecoder, PromptEncoder, Sam, TwoWayTransformer 11 | 12 | 13 | def build_sam_vit_h(image_size, num_classes, pixel_mean=[123.675, 116.28, 103.53], pixel_std=[58.395, 57.12, 57.375], 14 | checkpoint=None): 15 | return _build_sam( 16 | encoder_embed_dim=1280, 17 | encoder_depth=32, 18 | encoder_num_heads=16, 19 | encoder_global_attn_indexes=[7, 15, 23, 31], 20 | checkpoint=checkpoint, 21 | num_classes=num_classes, 22 | image_size=image_size, 23 | pixel_mean=pixel_mean, 24 | pixel_std=pixel_std 25 | ) 26 | 27 | 28 | build_sam = build_sam_vit_h 29 | 30 | 31 | def build_sam_vit_l(image_size, num_classes, pixel_mean=[123.675, 116.28, 103.53], pixel_std=[58.395, 57.12, 57.375], 32 | checkpoint=None): 33 | return _build_sam( 34 | encoder_embed_dim=1024, 35 | encoder_depth=24, 36 | encoder_num_heads=16, 37 | encoder_global_attn_indexes=[5, 11, 17, 23], 38 | checkpoint=checkpoint, 39 | num_classes=num_classes, 40 | image_size=image_size, 41 | pixel_mean=pixel_mean, 42 | pixel_std=pixel_std 43 | ) 44 | 45 | 46 | def build_sam_vit_b(image_size, num_classes, pixel_mean=[123.675, 116.28, 103.53], pixel_std=[58.395, 57.12, 57.375], 47 | checkpoint=None): 48 | return _build_sam( 49 | encoder_embed_dim=768, 50 | encoder_depth=12, 51 | encoder_num_heads=12, 52 | encoder_global_attn_indexes=[2, 5, 8, 11], 53 | # adopt global attention at [3, 6, 9, 12] transform layer, else window attention layer 54 | checkpoint=checkpoint, 55 | num_classes=num_classes, 56 | image_size=image_size, 57 | pixel_mean=pixel_mean, 58 | pixel_std=pixel_std 59 | ) 60 | 61 | 62 | sam_model_registry = { 63 | "default": build_sam_vit_h, 64 | "vit_h": build_sam_vit_h, 65 | "vit_l": build_sam_vit_l, 66 | "vit_b": build_sam_vit_b, 67 | } 68 | 69 | 70 | def _build_sam( 71 | encoder_embed_dim, 72 | encoder_depth, 73 | encoder_num_heads, 74 | encoder_global_attn_indexes, 75 | num_classes, 76 | image_size, 77 | pixel_mean, 78 | pixel_std, 79 | checkpoint=None, 80 | ): 81 | prompt_embed_dim = 256 82 | image_size = image_size 83 | vit_patch_size = 16 84 | image_embedding_size = image_size // vit_patch_size # Divide by 16 here 85 | sam = Sam( 86 | image_encoder=ImageEncoderViT( 87 | depth=encoder_depth, 88 | embed_dim=encoder_embed_dim, 89 | img_size=image_size, 90 | mlp_ratio=4, 91 | norm_layer=partial(torch.nn.LayerNorm, eps=1e-6), 92 | num_heads=encoder_num_heads, 93 | patch_size=vit_patch_size, 94 | qkv_bias=True, 95 | use_rel_pos=True, 96 | global_attn_indexes=encoder_global_attn_indexes, 97 | window_size=14, 98 | out_chans=prompt_embed_dim, 99 | ), 100 | prompt_encoder=PromptEncoder( 101 | embed_dim=prompt_embed_dim, 102 | image_embedding_size=(image_embedding_size, image_embedding_size), 103 | input_image_size=(image_size, image_size), 104 | mask_in_chans=16, 105 | ), 106 | mask_decoder=MaskDecoder( 107 | num_multimask_outputs=num_classes, 108 | transformer=TwoWayTransformer( 109 | depth=2, 110 | embedding_dim=prompt_embed_dim, 111 | mlp_dim=2048, 112 | num_heads=8, 113 | ), 114 | transformer_dim=prompt_embed_dim, 115 | iou_head_depth=3, 116 | iou_head_hidden_dim=256, 117 | ), 118 | # pixel_mean=[123.675, 116.28, 103.53], 119 | # pixel_std=[58.395, 57.12, 57.375], 120 | pixel_mean=pixel_mean, 121 | pixel_std=pixel_std 122 | ) 123 | # sam.eval() 124 | sam.train() 125 | if checkpoint is not None: 126 | with open(checkpoint, "rb") as f: 127 | state_dict = torch.load(f) 128 | try: 129 | sam.load_state_dict(state_dict) 130 | except: 131 | new_state_dict = load_from(sam, state_dict, image_size, vit_patch_size) 132 | sam.load_state_dict(new_state_dict) 133 | return sam, image_embedding_size 134 | 135 | 136 | def load_from(sam, state_dict, image_size, vit_patch_size): 137 | sam_dict = sam.state_dict() 138 | except_keys = ['mask_tokens', 'output_hypernetworks_mlps', 'iou_prediction_head'] 139 | new_state_dict = {k: v for k, v in state_dict.items() if 140 | k in sam_dict.keys() and except_keys[0] not in k and except_keys[1] not in k and except_keys[2] not in k} 141 | pos_embed = new_state_dict['image_encoder.pos_embed'] 142 | token_size = int(image_size // vit_patch_size) 143 | if pos_embed.shape[1] != token_size: 144 | # resize pos embedding, which may sacrifice the performance, but I have no better idea 145 | pos_embed = pos_embed.permute(0, 3, 1, 2) # [b, c, h, w] # [1, 768, 64, 64] 146 | pos_embed = F.interpolate(pos_embed, (token_size, token_size), mode='bilinear', align_corners=False) # [1, 768, 32, 32] 147 | pos_embed = pos_embed.permute(0, 2, 3, 1) # [b, h, w, c] # [1, 32, 32, 768] 148 | new_state_dict['image_encoder.pos_embed'] = pos_embed 149 | rel_pos_keys = [k for k in sam_dict.keys() if 'rel_pos' in k] 150 | global_rel_pos_keys = [k for k in rel_pos_keys if '2' in k or '5' in k or '8' in k or '11' in k] 151 | for k in global_rel_pos_keys: 152 | rel_pos_params = new_state_dict[k] 153 | h, w = rel_pos_params.shape 154 | rel_pos_params = rel_pos_params.unsqueeze(0).unsqueeze(0) 155 | rel_pos_params = F.interpolate(rel_pos_params, (token_size * 2 - 1, w), mode='bilinear', align_corners=False) 156 | new_state_dict[k] = rel_pos_params[0, 0, ...] 157 | sam_dict.update(new_state_dict) 158 | return sam_dict 159 | 160 | -------------------------------------------------------------------------------- /code/segment_anything_lora/modeling/mask_decoder.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import torch 8 | from torch import nn 9 | from torch.nn import functional as F 10 | # from icecream import ic 11 | 12 | from typing import List, Tuple, Type 13 | 14 | from .common import LayerNorm2d 15 | 16 | 17 | class MaskDecoder(nn.Module): 18 | def __init__( 19 | self, 20 | *, 21 | transformer_dim: int, 22 | transformer: nn.Module, 23 | num_multimask_outputs: int = 3, 24 | activation: Type[nn.Module] = nn.GELU, 25 | iou_head_depth: int = 3, 26 | iou_head_hidden_dim: int = 256, 27 | ) -> None: 28 | """ 29 | Predicts masks given an image and prompt embeddings, using a 30 | tranformer architecture. 31 | 32 | Arguments: 33 | transformer_dim (int): the channel dimension of the transformer 34 | transformer (nn.Module): the transformer used to predict masks 35 | num_multimask_outputs (int): the number of masks to predict 36 | when disambiguating masks 37 | activation (nn.Module): the type of activation to use when 38 | upscaling masks 39 | iou_head_depth (int): the depth of the MLP used to predict 40 | mask quality 41 | iou_head_hidden_dim (int): the hidden dimension of the MLP 42 | used to predict mask quality 43 | """ 44 | super().__init__() 45 | self.transformer_dim = transformer_dim 46 | self.transformer = transformer 47 | 48 | self.num_multimask_outputs = num_multimask_outputs 49 | 50 | self.iou_token = nn.Embedding(1, transformer_dim) 51 | self.num_mask_tokens = num_multimask_outputs + 1 52 | self.mask_tokens = nn.Embedding(self.num_mask_tokens, transformer_dim) 53 | 54 | self.output_upscaling = nn.Sequential( 55 | nn.ConvTranspose2d(transformer_dim, transformer_dim // 4, kernel_size=2, stride=2), 56 | LayerNorm2d(transformer_dim // 4), 57 | activation(), 58 | nn.ConvTranspose2d(transformer_dim // 4, transformer_dim // 8, kernel_size=2, stride=2), 59 | activation(), 60 | ) 61 | self.output_hypernetworks_mlps = nn.ModuleList( 62 | [ 63 | MLP(transformer_dim, transformer_dim, transformer_dim // 8, 3) 64 | for i in range(self.num_mask_tokens) 65 | ] 66 | ) 67 | 68 | self.iou_prediction_head = MLP( 69 | transformer_dim, iou_head_hidden_dim, self.num_mask_tokens, iou_head_depth 70 | ) 71 | 72 | def forward( 73 | self, 74 | image_embeddings: torch.Tensor, 75 | image_pe: torch.Tensor, 76 | sparse_prompt_embeddings: torch.Tensor, 77 | dense_prompt_embeddings: torch.Tensor, 78 | multimask_output: bool, 79 | ) -> Tuple[torch.Tensor, torch.Tensor]: 80 | """ 81 | Predict masks given image and prompt embeddings. 82 | 83 | Arguments: 84 | image_embeddings (torch.Tensor): the embeddings from the image encoder 85 | image_pe (torch.Tensor): positional encoding with the shape of image_embeddings 86 | sparse_prompt_embeddings (torch.Tensor): the embeddings of the points and boxes 87 | dense_prompt_embeddings (torch.Tensor): the embeddings of the mask inputs 88 | multimask_output (bool): Whether to return multiple masks or a single 89 | mask. 90 | 91 | Returns: 92 | torch.Tensor: batched predicted masks 93 | torch.Tensor: batched predictions of mask quality 94 | """ 95 | masks, iou_pred = self.predict_masks( 96 | image_embeddings=image_embeddings, 97 | image_pe=image_pe, 98 | sparse_prompt_embeddings=sparse_prompt_embeddings, 99 | dense_prompt_embeddings=dense_prompt_embeddings, 100 | ) 101 | 102 | # Select the correct mask or masks for output 103 | # if multimask_output: 104 | # mask_slice = slice(1, None) 105 | # else: 106 | # mask_slice = slice(0, 1) 107 | # masks = masks[:, mask_slice, :, :] 108 | # iou_pred = iou_pred[:, mask_slice] 109 | 110 | # Prepare output 111 | return masks, iou_pred 112 | 113 | def predict_masks( 114 | self, 115 | image_embeddings: torch.Tensor, 116 | image_pe: torch.Tensor, 117 | sparse_prompt_embeddings: torch.Tensor, 118 | dense_prompt_embeddings: torch.Tensor, 119 | ) -> Tuple[torch.Tensor, torch.Tensor]: 120 | """Predicts masks. See 'forward' for more details.""" 121 | # Concatenate output tokens 122 | output_tokens = torch.cat([self.iou_token.weight, self.mask_tokens.weight], dim=0) 123 | output_tokens = output_tokens.unsqueeze(0).expand(sparse_prompt_embeddings.size(0), -1, -1) 124 | tokens = torch.cat((output_tokens, sparse_prompt_embeddings), dim=1) 125 | 126 | # Expand per-image data in batch direction to be per-mask 127 | src = torch.repeat_interleave(image_embeddings, tokens.shape[0], dim=0) 128 | src = src + dense_prompt_embeddings 129 | pos_src = torch.repeat_interleave(image_pe, tokens.shape[0], dim=0) 130 | b, c, h, w = src.shape 131 | 132 | # Run the transformer 133 | hs, src = self.transformer(src, pos_src, tokens) 134 | iou_token_out = hs[:, 0, :] 135 | mask_tokens_out = hs[:, 1 : (1 + self.num_mask_tokens), :] 136 | 137 | # Upscale mask embeddings and predict masks using the mask tokens 138 | src = src.transpose(1, 2).view(b, c, h, w) 139 | upscaled_embedding = self.output_upscaling(src) 140 | hyper_in_list: List[torch.Tensor] = [] 141 | for i in range(self.num_mask_tokens): 142 | hyper_in_list.append(self.output_hypernetworks_mlps[i](mask_tokens_out[:, i, :])) 143 | hyper_in = torch.stack(hyper_in_list, dim=1) # [b, c, token_num] 144 | 145 | b, c, h, w = upscaled_embedding.shape # [h, token_num, h, w] 146 | masks = (hyper_in @ upscaled_embedding.view(b, c, h * w)).view(b, -1, h, w) # [1, 4, 256, 256], 256 = 4 * 64, the size of image embeddings 147 | 148 | # Generate mask quality predictions 149 | iou_pred = self.iou_prediction_head(iou_token_out) 150 | 151 | return masks, iou_pred 152 | 153 | 154 | # Lightly adapted from 155 | # https://github.com/facebookresearch/MaskFormer/blob/main/mask_former/modeling/transformer/transformer_predictor.py # noqa 156 | class MLP(nn.Module): 157 | def __init__( 158 | self, 159 | input_dim: int, 160 | hidden_dim: int, 161 | output_dim: int, 162 | num_layers: int, 163 | sigmoid_output: bool = False, 164 | ) -> None: 165 | super().__init__() 166 | self.num_layers = num_layers 167 | h = [hidden_dim] * (num_layers - 1) 168 | self.layers = nn.ModuleList( 169 | nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]) 170 | ) 171 | self.sigmoid_output = sigmoid_output 172 | 173 | def forward(self, x): 174 | for i, layer in enumerate(self.layers): 175 | x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x) 176 | if self.sigmoid_output: 177 | x = F.sigmoid(x) 178 | return x 179 | -------------------------------------------------------------------------------- /code/sam_lora_image_encoder.py: -------------------------------------------------------------------------------- 1 | from segment_anything import build_sam, SamPredictor 2 | from segment_anything import sam_model_registry 3 | 4 | import math 5 | import torch 6 | import torch.nn as nn 7 | import torch.nn.functional as F 8 | from torch import Tensor 9 | from torch.nn.parameter import Parameter 10 | from segment_anything.modeling import Sam 11 | # from safetensors import safe_open 12 | # from safetensors.torch import save_file 13 | 14 | # from icecream import ic 15 | 16 | 17 | class _LoRA_qkv(nn.Module): 18 | """In Sam it is implemented as 19 | self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) 20 | B, N, C = x.shape 21 | qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4) 22 | q, k, v = qkv.unbind(0) 23 | """ 24 | 25 | def __init__( 26 | self, 27 | qkv: nn.Module, 28 | linear_a_q: nn.Module, 29 | linear_b_q: nn.Module, 30 | linear_a_v: nn.Module, 31 | linear_b_v: nn.Module, 32 | ): 33 | super().__init__() 34 | self.qkv = qkv 35 | self.linear_a_q = linear_a_q 36 | self.linear_b_q = linear_b_q 37 | self.linear_a_v = linear_a_v 38 | self.linear_b_v = linear_b_v 39 | self.dim = qkv.in_features 40 | self.w_identity = torch.eye(qkv.in_features) 41 | 42 | def forward(self, x): 43 | qkv = self.qkv(x) # B,N,N,3*org_C 44 | new_q = self.linear_b_q(self.linear_a_q(x)) 45 | new_v = self.linear_b_v(self.linear_a_v(x)) 46 | qkv[:, :, :, : self.dim] += new_q 47 | qkv[:, :, :, -self.dim:] += new_v 48 | return qkv 49 | 50 | 51 | class LoRA_Sam(nn.Module): 52 | """Applies low-rank adaptation to a Sam model's image encoder. 53 | 54 | Args: 55 | sam_model: a vision transformer model, see base_vit.py 56 | r: rank of LoRA 57 | num_classes: how many classes the model output, default to the vit model 58 | lora_layer: which layer we apply LoRA. 59 | 60 | Examples:: 61 | >>> model = ViT('B_16_imagenet1k') 62 | >>> lora_model = LoRA_ViT(model, r=4) 63 | >>> preds = lora_model(img) 64 | >>> print(preds.shape) 65 | torch.Size([1, 1000]) 66 | """ 67 | 68 | def __init__(self, sam_model: Sam, r: int, lora_layer=None): 69 | super(LoRA_Sam, self).__init__() 70 | 71 | assert r > 0 72 | # base_vit_dim = sam_model.image_encoder.patch_embed.proj.out_channels 73 | # dim = base_vit_dim 74 | if lora_layer: 75 | self.lora_layer = lora_layer 76 | else: 77 | self.lora_layer = list( 78 | range(len(sam_model.image_encoder.blocks))) # Only apply lora to the image encoder by default 79 | # create for storage, then we can init them or load weights 80 | self.w_As = [] # These are linear layers 81 | self.w_Bs = [] 82 | 83 | # lets freeze first 84 | for param in sam_model.image_encoder.parameters(): 85 | param.requires_grad = False 86 | 87 | # Here, we do the surgery 88 | for t_layer_i, blk in enumerate(sam_model.image_encoder.blocks): 89 | # If we only want few lora layer instead of all 90 | if t_layer_i not in self.lora_layer: 91 | continue 92 | w_qkv_linear = blk.attn.qkv 93 | self.dim = w_qkv_linear.in_features 94 | w_a_linear_q = nn.Linear(self.dim, r, bias=False) 95 | w_b_linear_q = nn.Linear(r, self.dim, bias=False) 96 | w_a_linear_v = nn.Linear(self.dim, r, bias=False) 97 | w_b_linear_v = nn.Linear(r, self.dim, bias=False) 98 | self.w_As.append(w_a_linear_q) 99 | self.w_Bs.append(w_b_linear_q) 100 | self.w_As.append(w_a_linear_v) 101 | self.w_Bs.append(w_b_linear_v) 102 | blk.attn.qkv = _LoRA_qkv( 103 | w_qkv_linear, 104 | w_a_linear_q, 105 | w_b_linear_q, 106 | w_a_linear_v, 107 | w_b_linear_v, 108 | ) 109 | self.reset_parameters() 110 | self.sam = sam_model 111 | 112 | def save_lora_parameters(self, filename: str) -> None: 113 | r"""Only safetensors is supported now. 114 | 115 | pip install safetensor if you do not have one installed yet. 116 | 117 | save both lora and fc parameters. 118 | """ 119 | 120 | assert filename.endswith(".pt") or filename.endswith('.pth') 121 | 122 | num_layer = len(self.w_As) # actually, it is half 123 | a_tensors = {f"w_a_{i:03d}": self.w_As[i].weight for i in range(num_layer)} 124 | b_tensors = {f"w_b_{i:03d}": self.w_Bs[i].weight for i in range(num_layer)} 125 | prompt_encoder_tensors = {} 126 | mask_decoder_tensors = {} 127 | 128 | # save prompt encoder, only `state_dict`, the `named_parameter` is not permitted 129 | if isinstance(self.sam, torch.nn.DataParallel) or isinstance(self.sam, torch.nn.parallel.DistributedDataParallel): 130 | state_dict = self.sam.module.state_dict() 131 | else: 132 | state_dict = self.sam.state_dict() 133 | for key, value in state_dict.items(): 134 | if 'prompt_encoder' in key: 135 | prompt_encoder_tensors[key] = value 136 | if 'mask_decoder' in key: 137 | mask_decoder_tensors[key] = value 138 | 139 | merged_dict = {**a_tensors, **b_tensors, **prompt_encoder_tensors, **mask_decoder_tensors} 140 | torch.save(merged_dict, filename) 141 | 142 | def load_lora_parameters(self, filename: str) -> None: 143 | r"""Only safetensors is supported now. 144 | 145 | pip install safetensor if you do not have one installed yet.\ 146 | 147 | load both lora and fc parameters. 148 | """ 149 | 150 | assert filename.endswith(".pt") or filename.endswith('.pth') 151 | 152 | state_dict = torch.load(filename) 153 | 154 | for i, w_A_linear in enumerate(self.w_As): 155 | saved_key = f"w_a_{i:03d}" 156 | saved_tensor = state_dict[saved_key] 157 | w_A_linear.weight = Parameter(saved_tensor) 158 | 159 | for i, w_B_linear in enumerate(self.w_Bs): 160 | saved_key = f"w_b_{i:03d}" 161 | saved_tensor = state_dict[saved_key] 162 | w_B_linear.weight = Parameter(saved_tensor) 163 | 164 | sam_dict = self.sam.state_dict() 165 | sam_keys = sam_dict.keys() 166 | 167 | # load prompt encoder 168 | prompt_encoder_keys = [k for k in sam_keys if 'prompt_encoder' in k] 169 | prompt_encoder_values = [state_dict[k] for k in prompt_encoder_keys] 170 | prompt_encoder_new_state_dict = {k: v for k, v in zip(prompt_encoder_keys, prompt_encoder_values)} 171 | sam_dict.update(prompt_encoder_new_state_dict) 172 | 173 | # load mask decoder 174 | mask_decoder_keys = [k for k in sam_keys if 'mask_decoder' in k] 175 | mask_decoder_values = [state_dict[k] for k in mask_decoder_keys] 176 | mask_decoder_new_state_dict = {k: v for k, v in zip(mask_decoder_keys, mask_decoder_values)} 177 | sam_dict.update(mask_decoder_new_state_dict) 178 | self.sam.load_state_dict(sam_dict) 179 | 180 | def reset_parameters(self) -> None: 181 | for w_A in self.w_As: 182 | nn.init.kaiming_uniform_(w_A.weight, a=math.sqrt(5)) 183 | for w_B in self.w_Bs: 184 | nn.init.zeros_(w_B.weight) 185 | 186 | def forward(self, batched_input, multimask_output, image_size): 187 | return self.sam(batched_input, multimask_output, image_size) 188 | 189 | 190 | # def forward(self, x: Tensor) -> Tensor: 191 | # return self.lora_vit(x) 192 | 193 | 194 | if __name__ == "__main__": 195 | sam = sam_model_registry["vit_b"](checkpoint="sam_vit_b_01ec64.pth") 196 | lora_sam = LoRA_Sam(sam, 4) 197 | lora_sam.sam.image_encoder(torch.rand(size=(1, 3, 1024, 1024))) -------------------------------------------------------------------------------- /code/utils/losses.py: -------------------------------------------------------------------------------- 1 | import pdb 2 | 3 | import torch 4 | from torch.nn import functional as F 5 | import numpy as np 6 | import torch.nn as nn 7 | from torch.autograd import Variable 8 | import math 9 | import contextlib 10 | 11 | def dice_loss(score, target): 12 | target = target.float() 13 | smooth = 1e-5 14 | intersect = torch.sum(score * target) 15 | y_sum = torch.sum(target * target) 16 | z_sum = torch.sum(score * score) 17 | loss = (2 * intersect + smooth) / (z_sum + y_sum + smooth) 18 | loss = 1 - loss 19 | return loss 20 | 21 | def dice_loss1(score, target): 22 | target = target.float() 23 | smooth = 1e-5 24 | intersect = torch.sum(score * target) 25 | y_sum = torch.sum(target) 26 | z_sum = torch.sum(score) 27 | loss = (2 * intersect + smooth) / (z_sum + y_sum + smooth) 28 | loss = 1 - loss 29 | return loss 30 | 31 | def entropy_loss(p, C=2): 32 | # p N*C*W*H*D 33 | y1 = -1*torch.sum(p*torch.log(p+1e-6), dim=1) / torch.tensor(np.log(C)).cuda() 34 | ent = torch.mean(y1) 35 | 36 | return ent 37 | 38 | def softmax_dice_loss(input_logits, target_logits): 39 | """Takes softmax on both sides and returns MSE loss 40 | Note: 41 | - Returns the sum over all examples. Divide by the batch size afterwards 42 | if you want the mean. 43 | - Sends gradients to inputs but not the targets. 44 | """ 45 | assert input_logits.size() == target_logits.size() 46 | input_softmax = F.softmax(input_logits, dim=1) 47 | target_softmax = F.softmax(target_logits, dim=1) 48 | n = input_logits.shape[1] 49 | dice = 0 50 | for i in range(0, n): 51 | dice += dice_loss1(input_softmax[:, i], target_softmax[:, i]) 52 | mean_dice = dice / n 53 | 54 | return mean_dice 55 | 56 | 57 | def entropy_loss_map(p, C=2): 58 | ent = -1*torch.sum(p * torch.log(p + 1e-6), dim=1, keepdim=True)/torch.tensor(np.log(C)).cuda() 59 | return ent 60 | 61 | 62 | def mse_loss(input_logits, target_logits): 63 | assert input_logits.size() == target_logits.size() 64 | 65 | mse_loss = (input_logits-target_logits)**2 66 | return mse_loss 67 | 68 | 69 | def softmax_mse_loss(input_logits, target_logits, sigmoid=False): 70 | """Takes softmax on both sides and returns MSE loss 71 | Note: 72 | - Returns the sum over all examples. Divide by the batch size afterwards 73 | if you want the mean. 74 | - Sends gradients to inputs but not the targets. 75 | """ 76 | assert input_logits.size() == target_logits.size() 77 | if sigmoid: 78 | input_softmax = torch.sigmoid(input_logits) 79 | target_softmax = torch.sigmoid(target_logits) 80 | else: 81 | input_softmax = F.softmax(input_logits, dim=1) 82 | target_softmax = F.softmax(target_logits, dim=1) 83 | 84 | mse_loss = (input_softmax-target_softmax)**2 85 | return mse_loss 86 | 87 | 88 | def softmax_kl_loss(input_logits, target_logits, sigmoid=False): 89 | """Takes softmax on both sides and returns KL divergence 90 | Note: 91 | - Returns the sum over all examples. Divide by the batch size afterwards 92 | if you want the mean. 93 | - Sends gradients to inputs but not the targets. 94 | """ 95 | assert input_logits.size() == target_logits.size() 96 | if sigmoid: 97 | input_log_softmax = torch.log(torch.sigmoid(input_logits)) 98 | target_softmax = torch.sigmoid(target_logits) 99 | else: 100 | input_log_softmax = F.log_softmax(input_logits, dim=1) 101 | target_softmax = F.softmax(target_logits, dim=1) 102 | 103 | # return F.kl_div(input_log_softmax, target_softmax) 104 | kl_div = F.kl_div(input_log_softmax, target_softmax, reduction='mean') 105 | # mean_kl_div = torch.mean(0.2*kl_div[:,0,...]+0.8*kl_div[:,1,...]) 106 | return kl_div 107 | 108 | 109 | def symmetric_mse_loss(input1, input2): 110 | """Like F.mse_loss but sends gradients to both directions 111 | Note: 112 | - Returns the sum over all examples. Divide by the batch size afterwards 113 | if you want the mean. 114 | - Sends gradients to both input1 and input2. 115 | """ 116 | assert input1.size() == input2.size() 117 | return torch.mean((input1 - input2)**2) 118 | 119 | 120 | class FocalLoss(nn.Module): 121 | def __init__(self, gamma=2, alpha=None, size_average=True): 122 | super(FocalLoss, self).__init__() 123 | self.gamma = gamma 124 | self.alpha = alpha 125 | if isinstance(alpha, (float, int)): 126 | self.alpha = torch.Tensor([alpha, 1-alpha]) 127 | if isinstance(alpha, list): 128 | self.alpha = torch.Tensor(alpha) 129 | self.size_average = size_average 130 | 131 | def forward(self, input, target): 132 | if input.dim() > 2: 133 | # N,C,H,W => N,C,H*W 134 | input = input.view(input.size(0), input.size(1), -1) 135 | input = input.transpose(1, 2) # N,C,H*W => N,H*W,C 136 | input = input.contiguous().view(-1, input.size(2)) # N,H*W,C => N*H*W,C 137 | target = target.view(-1, 1) 138 | 139 | logpt = F.log_softmax(input, dim=1) 140 | logpt = logpt.gather(1, target) 141 | logpt = logpt.view(-1) 142 | pt = Variable(logpt.data.exp()) 143 | 144 | if self.alpha is not None: 145 | if self.alpha.type() != input.data.type(): 146 | self.alpha = self.alpha.type_as(input.data) 147 | at = self.alpha.gather(0, target.data.view(-1)) 148 | logpt = logpt * Variable(at) 149 | 150 | loss = -1 * (1-pt)**self.gamma * logpt 151 | if self.size_average: 152 | return loss.mean() 153 | else: 154 | return loss.sum() 155 | 156 | 157 | class DiceLoss(nn.Module): 158 | def __init__(self, n_classes): 159 | super(DiceLoss, self).__init__() 160 | self.n_classes = n_classes 161 | 162 | def _one_hot_encoder(self, input_tensor): 163 | tensor_list = [] 164 | for i in range(self.n_classes): 165 | temp_prob = input_tensor == i * torch.ones_like(input_tensor) 166 | tensor_list.append(temp_prob) 167 | output_tensor = torch.cat(tensor_list, dim=1) 168 | return output_tensor.float() 169 | 170 | def _dice_loss(self, score, target): 171 | target = target.float() 172 | smooth = 1e-10 173 | intersect = torch.sum(score * target) 174 | y_sum = torch.sum(target * target) 175 | z_sum = torch.sum(score * score) 176 | loss = (2 * intersect + smooth) / (z_sum + y_sum + smooth) 177 | loss = 1 - loss 178 | return loss 179 | 180 | def forward(self, inputs, target, weight=None, softmax=False): 181 | if softmax: 182 | inputs = torch.softmax(inputs, dim=1) 183 | target = self._one_hot_encoder(target) 184 | if weight is None: 185 | weight = [1] * self.n_classes 186 | assert inputs.size() == target.size(), 'predict & target shape do not match' 187 | class_wise_dice = [] 188 | loss = 0.0 189 | for i in range(0, self.n_classes): 190 | dice = self._dice_loss(inputs[:, i], target[:, i]) 191 | class_wise_dice.append(1.0 - dice.item()) 192 | loss += dice * weight[i] 193 | return loss / self.n_classes 194 | 195 | 196 | def to_one_hot(tensor, nClasses): 197 | """ Input tensor : Nx1xHxW 198 | :param tensor: 199 | :param nClasses: 200 | :return: 201 | """ 202 | assert tensor.max().item() < nClasses, 'one hot tensor.max() = {} < {}'.format(torch.max(tensor), nClasses) 203 | assert tensor.min().item() >= 0, 'one hot tensor.min() = {} < {}'.format(tensor.min(), 0) 204 | 205 | size = list(tensor.size()) 206 | assert size[1] == 1 207 | size[1] = nClasses 208 | one_hot = torch.zeros(*size) 209 | if tensor.is_cuda: 210 | one_hot = one_hot.cuda(tensor.device) 211 | one_hot = one_hot.scatter_(1, tensor, 1) 212 | return one_hot 213 | 214 | def get_probability(logits): 215 | """ Get probability from logits, if the channel of logits is 1 then use sigmoid else use softmax. 216 | :param logits: [N, C, H, W] or [N, C, D, H, W] 217 | :return: prediction and class num 218 | """ 219 | size = logits.size() 220 | # N x 1 x H x W 221 | if size[1] > 1: 222 | pred = F.softmax(logits, dim=1) 223 | nclass = size[1] 224 | else: 225 | pred = F.sigmoid(logits) 226 | pred = torch.cat([1 - pred, pred], 1) 227 | nclass = 2 228 | return pred, nclass 229 | 230 | 231 | def entropy_minmization(p): 232 | y1 = -1*torch.sum(p*torch.log(p+1e-6), dim=1) 233 | ent = torch.mean(y1) 234 | return ent 235 | 236 | 237 | def entropy_map(p): 238 | ent_map = -1*torch.sum(p * torch.log(p + 1e-6), dim=1, 239 | keepdim=True) 240 | return ent_map 241 | 242 | 243 | 244 | def EMA(cur_weight, past_weight, momentum=0.9): 245 | new_weight = momentum * past_weight + (1 - momentum) * cur_weight 246 | return new_weight 247 | 248 | -------------------------------------------------------------------------------- /code/segment_anything_lora/modeling/sam.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import torch 8 | from torch import nn 9 | from torch.nn import functional as F 10 | # from icecream import ic 11 | 12 | from typing import Any, Dict, List, Tuple 13 | 14 | from .image_encoder import ImageEncoderViT 15 | from .mask_decoder import MaskDecoder 16 | from .prompt_encoder import PromptEncoder 17 | 18 | 19 | class Sam(nn.Module): 20 | mask_threshold: float = 0.0 21 | image_format: str = "RGB" 22 | 23 | def __init__( 24 | self, 25 | image_encoder: ImageEncoderViT, 26 | prompt_encoder: PromptEncoder, 27 | mask_decoder: MaskDecoder, 28 | pixel_mean: List[float] = [123.675, 116.28, 103.53], 29 | pixel_std: List[float] = [58.395, 57.12, 57.375], 30 | ) -> None: 31 | """ 32 | SAM predicts object masks from an image and input prompts. 33 | 34 | Arguments: 35 | image_encoder (ImageEncoderViT): The backbone used to encode the 36 | image into image embeddings that allow for efficient mask prediction. 37 | prompt_encoder (PromptEncoder): Encodes various types of input prompts. 38 | mask_decoder (MaskDecoder): Predicts masks from the image embeddings 39 | and encoded prompts. 40 | pixel_mean (list(float)): Mean values for normalizing pixels in the input image. 41 | pixel_std (list(float)): Std values for normalizing pixels in the input image. 42 | """ 43 | super().__init__() 44 | self.image_encoder = image_encoder 45 | self.prompt_encoder = prompt_encoder 46 | self.mask_decoder = mask_decoder 47 | self.register_buffer("pixel_mean", torch.Tensor(pixel_mean).view(-1, 1, 1), False) 48 | self.register_buffer("pixel_std", torch.Tensor(pixel_std).view(-1, 1, 1), False) 49 | 50 | @property 51 | def device(self) -> Any: 52 | return self.pixel_mean.device 53 | 54 | def forward(self, batched_input, multimask_output, image_size): 55 | if isinstance(batched_input, list): 56 | outputs = self.forward_test(batched_input, multimask_output) 57 | else: 58 | outputs = self.forward_train(batched_input, multimask_output, image_size) 59 | return outputs 60 | 61 | def forward_train(self, batched_input, multimask_output, image_size): 62 | input_images = self.preprocess(batched_input) 63 | image_embeddings = self.image_encoder(input_images) 64 | sparse_embeddings, dense_embeddings = self.prompt_encoder( 65 | points=None, boxes=None, masks=None 66 | ) 67 | low_res_masks, iou_predictions = self.mask_decoder( 68 | image_embeddings=image_embeddings, 69 | image_pe=self.prompt_encoder.get_dense_pe(), 70 | sparse_prompt_embeddings=sparse_embeddings, 71 | dense_prompt_embeddings=dense_embeddings, 72 | multimask_output=multimask_output 73 | ) 74 | masks = self.postprocess_masks( 75 | low_res_masks, 76 | input_size=(image_size, image_size), 77 | original_size=(image_size, image_size) 78 | ) 79 | outputs = { 80 | 'masks': masks, 81 | 'iou_predictions': iou_predictions, 82 | 'low_res_logits': low_res_masks 83 | } 84 | return outputs 85 | 86 | @torch.no_grad() 87 | def forward_test( 88 | self, 89 | batched_input: List[Dict[str, Any]], 90 | multimask_output: bool, 91 | ) -> List[Dict[str, torch.Tensor]]: 92 | """ 93 | Predicts masks end-to-end from provided images and prompts. 94 | If prompts are not known in advance, using SamPredictor is 95 | recommended over calling the model directly. 96 | 97 | Arguments: 98 | batched_input (list(dict)): A list over input images, each a 99 | dictionary with the following keys. A prompt key can be 100 | excluded if it is not present. 101 | 'image': The image as a torch tensor in 3xHxW format, 102 | already transformed for input to the model. 103 | 'original_size': (tuple(int, int)) The original size of 104 | the image before transformation, as (H, W). 105 | 'point_coords': (torch.Tensor) Batched point prompts for 106 | this image, with shape BxNx2. Already transformed to the 107 | input frame of the model. 108 | 'point_labels': (torch.Tensor) Batched labels for point prompts, 109 | with shape BxN. 110 | 'boxes': (torch.Tensor) Batched box inputs, with shape Bx4. 111 | Already transformed to the input frame of the model. 112 | 'mask_inputs': (torch.Tensor) Batched mask inputs to the model, 113 | in the form Bx1xHxW. 114 | multimask_output (bool): Whether the model should predict multiple 115 | disambiguating masks, or return a single mask. 116 | 117 | Returns: 118 | (list(dict)): A list over input images, where each element is 119 | as dictionary with the following keys. 120 | 'masks': (torch.Tensor) Batched binary mask predictions, 121 | with shape BxCxHxW, where B is the number of input promts, 122 | C is determiend by multimask_output, and (H, W) is the 123 | original size of the image. 124 | 'iou_predictions': (torch.Tensor) The model's predictions 125 | of mask quality, in shape BxC. 126 | 'low_res_logits': (torch.Tensor) Low resolution logits with 127 | shape BxCxHxW, where H=W=256. Can be passed as mask input 128 | to subsequent iterations of prediction. 129 | """ 130 | input_images = torch.stack([self.preprocess(x["image"]) for x in batched_input], dim=0) 131 | image_embeddings = self.image_encoder(input_images) 132 | 133 | outputs = [] 134 | for image_record, curr_embedding in zip(batched_input, image_embeddings): 135 | if "point_coords" in image_record: 136 | points = (image_record["point_coords"], image_record["point_labels"]) 137 | else: 138 | points = None 139 | sparse_embeddings, dense_embeddings = self.prompt_encoder( 140 | points=points, 141 | boxes=image_record.get("boxes", None), 142 | masks=image_record.get("mask_inputs", None), 143 | ) 144 | low_res_masks, iou_predictions = self.mask_decoder( 145 | image_embeddings=curr_embedding.unsqueeze(0), 146 | image_pe=self.prompt_encoder.get_dense_pe(), 147 | sparse_prompt_embeddings=sparse_embeddings, 148 | dense_prompt_embeddings=dense_embeddings, 149 | multimask_output=multimask_output, 150 | ) 151 | masks = self.postprocess_masks( 152 | low_res_masks, 153 | input_size=image_record["image"].shape[-2:], 154 | original_size=image_record["original_size"], 155 | ) 156 | masks = masks > self.mask_threshold 157 | outputs.append( 158 | { 159 | "masks": masks, 160 | "iou_predictions": iou_predictions, 161 | "low_res_logits": low_res_masks, 162 | } 163 | ) 164 | return outputs 165 | 166 | def postprocess_masks( 167 | self, 168 | masks: torch.Tensor, 169 | input_size: Tuple[int, ...], 170 | original_size: Tuple[int, ...], 171 | ) -> torch.Tensor: 172 | """ 173 | Remove padding and upscale masks to the original image size. 174 | 175 | Arguments: 176 | masks (torch.Tensor): Batched masks from the mask_decoder, 177 | in BxCxHxW format. 178 | input_size (tuple(int, int)): The size of the image input to the 179 | model, in (H, W) format. Used to remove padding. 180 | original_size (tuple(int, int)): The original size of the image 181 | before resizing for input to the model, in (H, W) format. 182 | 183 | Returns: 184 | (torch.Tensor): Batched masks in BxCxHxW format, where (H, W) 185 | is given by original_size. 186 | """ 187 | masks = F.interpolate( 188 | masks, 189 | (self.image_encoder.img_size, self.image_encoder.img_size), 190 | mode="bilinear", 191 | align_corners=False, 192 | ) 193 | masks = masks[..., : input_size[0], : input_size[1]] 194 | masks = F.interpolate(masks, original_size, mode="bilinear", align_corners=False) 195 | return masks 196 | 197 | def preprocess(self, x: torch.Tensor) -> torch.Tensor: 198 | """Normalize pixel values and pad to a square input.""" 199 | # Normalize colors 200 | x = (x - self.pixel_mean) / self.pixel_std 201 | 202 | # Pad 203 | h, w = x.shape[-2:] 204 | padh = self.image_encoder.img_size - h 205 | padw = self.image_encoder.img_size - w 206 | x = F.pad(x, (0, padw, 0, padh)) 207 | return x 208 | 209 | -------------------------------------------------------------------------------- /code/segment_anything_lora/modeling/transformer.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import torch 8 | from torch import Tensor, nn 9 | 10 | import math 11 | from typing import Tuple, Type 12 | 13 | from .common import MLPBlock 14 | 15 | 16 | class TwoWayTransformer(nn.Module): 17 | def __init__( 18 | self, 19 | depth: int, 20 | embedding_dim: int, 21 | num_heads: int, 22 | mlp_dim: int, 23 | activation: Type[nn.Module] = nn.ReLU, 24 | attention_downsample_rate: int = 2, 25 | ) -> None: 26 | """ 27 | A transformer decoder that attends to an input image using 28 | queries whose positional embedding is supplied. 29 | 30 | Args: 31 | depth (int): number of layers in the transformer 32 | embedding_dim (int): the channel dimension for the input embeddings 33 | num_heads (int): the number of heads for multihead attention. Must 34 | divide embedding_dim 35 | mlp_dim (int): the channel dimension internal to the MLP block 36 | activation (nn.Module): the activation to use in the MLP block 37 | """ 38 | super().__init__() 39 | self.depth = depth 40 | self.embedding_dim = embedding_dim 41 | self.num_heads = num_heads 42 | self.mlp_dim = mlp_dim 43 | self.layers = nn.ModuleList() 44 | 45 | for i in range(depth): 46 | self.layers.append( 47 | TwoWayAttentionBlock( 48 | embedding_dim=embedding_dim, 49 | num_heads=num_heads, 50 | mlp_dim=mlp_dim, 51 | activation=activation, 52 | attention_downsample_rate=attention_downsample_rate, 53 | skip_first_layer_pe=(i == 0), 54 | ) 55 | ) 56 | 57 | self.final_attn_token_to_image = Attention( 58 | embedding_dim, num_heads, downsample_rate=attention_downsample_rate 59 | ) 60 | self.norm_final_attn = nn.LayerNorm(embedding_dim) 61 | 62 | def forward( 63 | self, 64 | image_embedding: Tensor, 65 | image_pe: Tensor, 66 | point_embedding: Tensor, 67 | ) -> Tuple[Tensor, Tensor]: 68 | """ 69 | Args: 70 | image_embedding (torch.Tensor): image to attend to. Should be shape 71 | B x embedding_dim x h x w for any h and w. 72 | image_pe (torch.Tensor): the positional encoding to add to the image. Must 73 | have the same shape as image_embedding. 74 | point_embedding (torch.Tensor): the embedding to add to the query points. 75 | Must have shape B x N_points x embedding_dim for any N_points. 76 | 77 | Returns: 78 | torch.Tensor: the processed point_embedding 79 | torch.Tensor: the processed image_embedding 80 | """ 81 | # BxCxHxW -> BxHWxC == B x N_image_tokens x C 82 | bs, c, h, w = image_embedding.shape 83 | image_embedding = image_embedding.flatten(2).permute(0, 2, 1) 84 | image_pe = image_pe.flatten(2).permute(0, 2, 1) 85 | 86 | # Prepare queries 87 | queries = point_embedding 88 | keys = image_embedding 89 | 90 | # Apply transformer blocks and final layernorm 91 | for layer in self.layers: 92 | queries, keys = layer( 93 | queries=queries, 94 | keys=keys, 95 | query_pe=point_embedding, 96 | key_pe=image_pe, 97 | ) 98 | 99 | # Apply the final attenion layer from the points to the image 100 | q = queries + point_embedding 101 | k = keys + image_pe 102 | attn_out = self.final_attn_token_to_image(q=q, k=k, v=keys) 103 | queries = queries + attn_out 104 | queries = self.norm_final_attn(queries) 105 | 106 | return queries, keys 107 | 108 | 109 | class TwoWayAttentionBlock(nn.Module): 110 | def __init__( 111 | self, 112 | embedding_dim: int, 113 | num_heads: int, 114 | mlp_dim: int = 2048, 115 | activation: Type[nn.Module] = nn.ReLU, 116 | attention_downsample_rate: int = 2, 117 | skip_first_layer_pe: bool = False, 118 | ) -> None: 119 | """ 120 | A transformer block with four layers: (1) self-attention of sparse 121 | inputs, (2) cross attention of sparse inputs to dense inputs, (3) mlp 122 | block on sparse inputs, and (4) cross attention of dense inputs to sparse 123 | inputs. 124 | 125 | Arguments: 126 | embedding_dim (int): the channel dimension of the embeddings 127 | num_heads (int): the number of heads in the attention layers 128 | mlp_dim (int): the hidden dimension of the mlp block 129 | activation (nn.Module): the activation of the mlp block 130 | skip_first_layer_pe (bool): skip the PE on the first layer 131 | """ 132 | super().__init__() 133 | self.self_attn = Attention(embedding_dim, num_heads) 134 | self.norm1 = nn.LayerNorm(embedding_dim) 135 | 136 | self.cross_attn_token_to_image = Attention( 137 | embedding_dim, num_heads, downsample_rate=attention_downsample_rate 138 | ) 139 | self.norm2 = nn.LayerNorm(embedding_dim) 140 | 141 | self.mlp = MLPBlock(embedding_dim, mlp_dim, activation) 142 | self.norm3 = nn.LayerNorm(embedding_dim) 143 | 144 | self.norm4 = nn.LayerNorm(embedding_dim) 145 | self.cross_attn_image_to_token = Attention( 146 | embedding_dim, num_heads, downsample_rate=attention_downsample_rate 147 | ) 148 | 149 | self.skip_first_layer_pe = skip_first_layer_pe 150 | 151 | def forward( 152 | self, queries: Tensor, keys: Tensor, query_pe: Tensor, key_pe: Tensor 153 | ) -> Tuple[Tensor, Tensor]: 154 | # Self attention block 155 | if self.skip_first_layer_pe: 156 | queries = self.self_attn(q=queries, k=queries, v=queries) 157 | else: 158 | q = queries + query_pe 159 | attn_out = self.self_attn(q=q, k=q, v=queries) 160 | queries = queries + attn_out 161 | queries = self.norm1(queries) 162 | 163 | # Cross attention block, tokens attending to image embedding 164 | q = queries + query_pe 165 | k = keys + key_pe 166 | attn_out = self.cross_attn_token_to_image(q=q, k=k, v=keys) 167 | queries = queries + attn_out 168 | queries = self.norm2(queries) 169 | 170 | # MLP block 171 | mlp_out = self.mlp(queries) 172 | queries = queries + mlp_out 173 | queries = self.norm3(queries) 174 | 175 | # Cross attention block, image embedding attending to tokens 176 | q = queries + query_pe 177 | k = keys + key_pe 178 | attn_out = self.cross_attn_image_to_token(q=k, k=q, v=queries) 179 | keys = keys + attn_out 180 | keys = self.norm4(keys) 181 | 182 | return queries, keys 183 | 184 | 185 | class Attention(nn.Module): 186 | """ 187 | An attention layer that allows for downscaling the size of the embedding 188 | after projection to queries, keys, and values. 189 | """ 190 | 191 | def __init__( 192 | self, 193 | embedding_dim: int, 194 | num_heads: int, 195 | downsample_rate: int = 1, 196 | ) -> None: 197 | super().__init__() 198 | self.embedding_dim = embedding_dim 199 | self.internal_dim = embedding_dim // downsample_rate 200 | self.num_heads = num_heads 201 | assert self.internal_dim % num_heads == 0, "num_heads must divide embedding_dim." 202 | 203 | self.q_proj = nn.Linear(embedding_dim, self.internal_dim) 204 | self.k_proj = nn.Linear(embedding_dim, self.internal_dim) 205 | self.v_proj = nn.Linear(embedding_dim, self.internal_dim) 206 | self.out_proj = nn.Linear(self.internal_dim, embedding_dim) 207 | 208 | def _separate_heads(self, x: Tensor, num_heads: int) -> Tensor: 209 | b, n, c = x.shape 210 | x = x.reshape(b, n, num_heads, c // num_heads) 211 | return x.transpose(1, 2) # B x N_heads x N_tokens x C_per_head 212 | 213 | def _recombine_heads(self, x: Tensor) -> Tensor: 214 | b, n_heads, n_tokens, c_per_head = x.shape 215 | x = x.transpose(1, 2) 216 | return x.reshape(b, n_tokens, n_heads * c_per_head) # B x N_tokens x C 217 | 218 | def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor: 219 | # Input projections 220 | q = self.q_proj(q) 221 | k = self.k_proj(k) 222 | v = self.v_proj(v) 223 | 224 | # Separate into heads 225 | q = self._separate_heads(q, self.num_heads) 226 | k = self._separate_heads(k, self.num_heads) 227 | v = self._separate_heads(v, self.num_heads) 228 | 229 | # Attention 230 | _, _, _, c_per_head = q.shape 231 | attn = q @ k.permute(0, 1, 3, 2) # B x N_heads x N_tokens x N_tokens 232 | attn = attn / math.sqrt(c_per_head) 233 | attn = torch.softmax(attn, dim=-1) 234 | 235 | # Get output 236 | out = attn @ v 237 | out = self._recombine_heads(out) 238 | out = self.out_proj(out) 239 | 240 | return out 241 | -------------------------------------------------------------------------------- /code/segment_anything_lora/modeling/prompt_encoder.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import numpy as np 8 | import torch 9 | from torch import nn 10 | 11 | from typing import Any, Optional, Tuple, Type 12 | 13 | from .common import LayerNorm2d 14 | 15 | 16 | class PromptEncoder(nn.Module): 17 | def __init__( 18 | self, 19 | embed_dim: int, 20 | image_embedding_size: Tuple[int, int], 21 | input_image_size: Tuple[int, int], 22 | mask_in_chans: int, 23 | activation: Type[nn.Module] = nn.GELU, 24 | ) -> None: 25 | """ 26 | Encodes prompts for input to SAM's mask decoder. 27 | 28 | Arguments: 29 | embed_dim (int): The prompts' embedding dimension 30 | image_embedding_size (tuple(int, int)): The spatial size of the 31 | image embedding, as (H, W). 32 | input_image_size (int): The padded size of the image as input 33 | to the image encoder, as (H, W). 34 | mask_in_chans (int): The number of hidden channels used for 35 | encoding input masks. 36 | activation (nn.Module): The activation to use when encoding 37 | input masks. 38 | """ 39 | super().__init__() 40 | self.embed_dim = embed_dim 41 | self.input_image_size = input_image_size 42 | self.image_embedding_size = image_embedding_size 43 | self.pe_layer = PositionEmbeddingRandom(embed_dim // 2) 44 | 45 | self.num_point_embeddings: int = 4 # pos/neg point + 2 box corners 46 | point_embeddings = [nn.Embedding(1, embed_dim) for i in range(self.num_point_embeddings)] 47 | self.point_embeddings = nn.ModuleList(point_embeddings) 48 | self.not_a_point_embed = nn.Embedding(1, embed_dim) 49 | 50 | self.mask_input_size = (4 * image_embedding_size[0], 4 * image_embedding_size[1]) 51 | self.mask_downscaling = nn.Sequential( 52 | nn.Conv2d(1, mask_in_chans // 4, kernel_size=2, stride=2), 53 | LayerNorm2d(mask_in_chans // 4), 54 | activation(), 55 | nn.Conv2d(mask_in_chans // 4, mask_in_chans, kernel_size=2, stride=2), 56 | LayerNorm2d(mask_in_chans), 57 | activation(), 58 | nn.Conv2d(mask_in_chans, embed_dim, kernel_size=1), 59 | ) # downsample to 1/4 60 | self.no_mask_embed = nn.Embedding(1, embed_dim) 61 | 62 | def get_dense_pe(self) -> torch.Tensor: 63 | """ 64 | Returns the positional encoding used to encode point prompts, 65 | applied to a dense set of points the shape of the image encoding. 66 | 67 | Returns: 68 | torch.Tensor: Positional encoding with shape 69 | 1x(embed_dim)x(embedding_h)x(embedding_w) 70 | """ 71 | return self.pe_layer(self.image_embedding_size).unsqueeze(0) 72 | 73 | def _embed_points( 74 | self, 75 | points: torch.Tensor, 76 | labels: torch.Tensor, 77 | pad: bool, 78 | ) -> torch.Tensor: 79 | """Embeds point prompts.""" 80 | points = points + 0.5 # Shift to center of pixel 81 | if pad: 82 | padding_point = torch.zeros((points.shape[0], 1, 2), device=points.device) 83 | padding_label = -torch.ones((labels.shape[0], 1), device=labels.device) 84 | points = torch.cat([points, padding_point], dim=1) 85 | labels = torch.cat([labels, padding_label], dim=1) 86 | point_embedding = self.pe_layer.forward_with_coords(points, self.input_image_size) 87 | point_embedding[labels == -1] = 0.0 88 | point_embedding[labels == -1] += self.not_a_point_embed.weight 89 | point_embedding[labels == 0] += self.point_embeddings[0].weight 90 | point_embedding[labels == 1] += self.point_embeddings[1].weight 91 | return point_embedding 92 | 93 | def _embed_boxes(self, boxes: torch.Tensor) -> torch.Tensor: 94 | """Embeds box prompts.""" 95 | boxes = boxes + 0.5 # Shift to center of pixel 96 | coords = boxes.reshape(-1, 2, 2) 97 | corner_embedding = self.pe_layer.forward_with_coords(coords, self.input_image_size) 98 | corner_embedding[:, 0, :] += self.point_embeddings[2].weight 99 | corner_embedding[:, 1, :] += self.point_embeddings[3].weight 100 | return corner_embedding 101 | 102 | def _embed_masks(self, masks: torch.Tensor) -> torch.Tensor: 103 | """Embeds mask inputs.""" 104 | mask_embedding = self.mask_downscaling(masks) 105 | return mask_embedding 106 | 107 | def _get_batch_size( 108 | self, 109 | points: Optional[Tuple[torch.Tensor, torch.Tensor]], 110 | boxes: Optional[torch.Tensor], 111 | masks: Optional[torch.Tensor], 112 | ) -> int: 113 | """ 114 | Gets the batch size of the output given the batch size of the input prompts. 115 | """ 116 | if points is not None: 117 | return points[0].shape[0] 118 | elif boxes is not None: 119 | return boxes.shape[0] 120 | elif masks is not None: 121 | return masks.shape[0] 122 | else: 123 | return 1 124 | 125 | def _get_device(self) -> torch.device: 126 | return self.point_embeddings[0].weight.device 127 | 128 | def forward( 129 | self, 130 | points: Optional[Tuple[torch.Tensor, torch.Tensor]], 131 | boxes: Optional[torch.Tensor], 132 | masks: Optional[torch.Tensor], 133 | ) -> Tuple[torch.Tensor, torch.Tensor]: 134 | """ 135 | Embeds different types of prompts, returning both sparse and dense 136 | embeddings. 137 | 138 | Arguments: 139 | points (tuple(torch.Tensor, torch.Tensor) or none): point coordinates 140 | and labels to embed. 141 | boxes (torch.Tensor or none): boxes to embed 142 | masks (torch.Tensor or none): masks to embed 143 | 144 | Returns: 145 | torch.Tensor: sparse embeddings for the points and boxes, with shape 146 | BxNx(embed_dim), where N is determined by the number of input points 147 | and boxes. 148 | torch.Tensor: dense embeddings for the masks, in the shape 149 | Bx(embed_dim)x(embed_H)x(embed_W) 150 | """ 151 | bs = self._get_batch_size(points, boxes, masks) 152 | sparse_embeddings = torch.empty((bs, 0, self.embed_dim), device=self._get_device()) 153 | if points is not None: 154 | coords, labels = points 155 | point_embeddings = self._embed_points(coords, labels, pad=(boxes is None)) 156 | sparse_embeddings = torch.cat([sparse_embeddings, point_embeddings], dim=1) 157 | if boxes is not None: 158 | box_embeddings = self._embed_boxes(boxes) 159 | sparse_embeddings = torch.cat([sparse_embeddings, box_embeddings], dim=1) 160 | 161 | if masks is not None: 162 | dense_embeddings = self._embed_masks(masks) 163 | else: 164 | dense_embeddings = self.no_mask_embed.weight.reshape(1, -1, 1, 1).expand( 165 | bs, -1, self.image_embedding_size[0], self.image_embedding_size[1] 166 | ) 167 | 168 | return sparse_embeddings, dense_embeddings 169 | 170 | 171 | class PositionEmbeddingRandom(nn.Module): 172 | """ 173 | Positional encoding using random spatial frequencies. 174 | """ 175 | 176 | def __init__(self, num_pos_feats: int = 64, scale: Optional[float] = None) -> None: 177 | super().__init__() 178 | if scale is None or scale <= 0.0: 179 | scale = 1.0 180 | self.register_buffer( 181 | "positional_encoding_gaussian_matrix", 182 | scale * torch.randn((2, num_pos_feats)), 183 | ) 184 | 185 | def _pe_encoding(self, coords: torch.Tensor) -> torch.Tensor: 186 | """Positionally encode points that are normalized to [0,1].""" 187 | # assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape 188 | coords = 2 * coords - 1 189 | coords = coords @ self.positional_encoding_gaussian_matrix 190 | coords = 2 * np.pi * coords 191 | # outputs d_1 x ... x d_n x C shape 192 | return torch.cat([torch.sin(coords), torch.cos(coords)], dim=-1) 193 | 194 | def forward(self, size: Tuple[int, int]) -> torch.Tensor: 195 | """Generate positional encoding for a grid of the specified size.""" 196 | h, w = size 197 | device: Any = self.positional_encoding_gaussian_matrix.device 198 | grid = torch.ones((h, w), device=device, dtype=torch.float32) 199 | y_embed = grid.cumsum(dim=0) - 0.5 200 | x_embed = grid.cumsum(dim=1) - 0.5 201 | y_embed = y_embed / h 202 | x_embed = x_embed / w 203 | 204 | pe = self._pe_encoding(torch.stack([x_embed, y_embed], dim=-1)) 205 | return pe.permute(2, 0, 1) # C x H x W 206 | 207 | def forward_with_coords( 208 | self, coords_input: torch.Tensor, image_size: Tuple[int, int] 209 | ) -> torch.Tensor: 210 | """Positionally encode points that are not normalized to [0,1].""" 211 | coords = coords_input.clone() 212 | coords[:, :, 0] = coords[:, :, 0] / image_size[1] 213 | coords[:, :, 1] = coords[:, :, 1] / image_size[0] 214 | return self._pe_encoding(coords.to(torch.float)) # B x N x C 215 | -------------------------------------------------------------------------------- /code/networks/hierarchical_vnet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | import torch.nn.functional as F 4 | 5 | class ConvBlock(nn.Module): 6 | def __init__(self, n_stages, n_filters_in, n_filters_out, normalization='none'): 7 | super(ConvBlock, self).__init__() 8 | 9 | ops = [] 10 | for i in range(n_stages): 11 | if i==0: 12 | input_channel = n_filters_in 13 | else: 14 | input_channel = n_filters_out 15 | 16 | ops.append(nn.Conv3d(input_channel, n_filters_out, 3, padding=1)) 17 | if normalization == 'batchnorm': 18 | ops.append(nn.BatchNorm3d(n_filters_out)) 19 | elif normalization == 'groupnorm': 20 | ops.append(nn.GroupNorm(num_groups=16, num_channels=n_filters_out)) 21 | elif normalization == 'instancenorm': 22 | ops.append(nn.InstanceNorm3d(n_filters_out)) 23 | elif normalization != 'none': 24 | assert False 25 | ops.append(nn.ReLU(inplace=True)) 26 | 27 | self.conv = nn.Sequential(*ops) 28 | 29 | def forward(self, x): 30 | x = self.conv(x) 31 | return x 32 | 33 | 34 | class ResidualConvBlock(nn.Module): 35 | def __init__(self, n_stages, n_filters_in, n_filters_out, normalization='none'): 36 | super(ResidualConvBlock, self).__init__() 37 | 38 | ops = [] 39 | for i in range(n_stages): 40 | if i == 0: 41 | input_channel = n_filters_in 42 | else: 43 | input_channel = n_filters_out 44 | 45 | ops.append(nn.Conv3d(input_channel, n_filters_out, 3, padding=1)) 46 | if normalization == 'batchnorm': 47 | ops.append(nn.BatchNorm3d(n_filters_out)) 48 | elif normalization == 'groupnorm': 49 | ops.append(nn.GroupNorm(num_groups=16, num_channels=n_filters_out)) 50 | elif normalization == 'instancenorm': 51 | ops.append(nn.InstanceNorm3d(n_filters_out)) 52 | elif normalization != 'none': 53 | assert False 54 | 55 | if i != n_stages-1: 56 | ops.append(nn.ReLU(inplace=True)) 57 | 58 | self.conv = nn.Sequential(*ops) 59 | self.relu = nn.ReLU(inplace=True) 60 | 61 | def forward(self, x): 62 | x = (self.conv(x) + x) 63 | x = self.relu(x) 64 | return x 65 | 66 | 67 | class DownsamplingConvBlock(nn.Module): 68 | def __init__(self, n_filters_in, n_filters_out, stride=2, normalization='none'): 69 | super(DownsamplingConvBlock, self).__init__() 70 | 71 | ops = [] 72 | if normalization != 'none': 73 | ops.append(nn.Conv3d(n_filters_in, n_filters_out, stride, padding=0, stride=stride)) 74 | if normalization == 'batchnorm': 75 | ops.append(nn.BatchNorm3d(n_filters_out)) 76 | elif normalization == 'groupnorm': 77 | ops.append(nn.GroupNorm(num_groups=16, num_channels=n_filters_out)) 78 | elif normalization == 'instancenorm': 79 | ops.append(nn.InstanceNorm3d(n_filters_out)) 80 | else: 81 | assert False 82 | else: 83 | ops.append(nn.Conv3d(n_filters_in, n_filters_out, stride, padding=0, stride=stride)) 84 | 85 | ops.append(nn.ReLU(inplace=True)) 86 | 87 | self.conv = nn.Sequential(*ops) 88 | 89 | def forward(self, x): 90 | x = self.conv(x) 91 | return x 92 | 93 | 94 | class UpsamplingDeconvBlock(nn.Module): 95 | def __init__(self, n_filters_in, n_filters_out, stride=2, normalization='none'): 96 | super(UpsamplingDeconvBlock, self).__init__() 97 | 98 | ops = [] 99 | if normalization != 'none': 100 | ops.append(nn.ConvTranspose3d(n_filters_in, n_filters_out, stride, padding=0, stride=stride)) 101 | if normalization == 'batchnorm': 102 | ops.append(nn.BatchNorm3d(n_filters_out)) 103 | elif normalization == 'groupnorm': 104 | ops.append(nn.GroupNorm(num_groups=16, num_channels=n_filters_out)) 105 | elif normalization == 'instancenorm': 106 | ops.append(nn.InstanceNorm3d(n_filters_out)) 107 | else: 108 | assert False 109 | else: 110 | ops.append(nn.ConvTranspose3d(n_filters_in, n_filters_out, stride, padding=0, stride=stride)) 111 | 112 | ops.append(nn.ReLU(inplace=True)) 113 | 114 | self.conv = nn.Sequential(*ops) 115 | 116 | def forward(self, x): 117 | x = self.conv(x) 118 | return x 119 | 120 | 121 | class Upsampling(nn.Module): 122 | def __init__(self, n_filters_in, n_filters_out, stride=2, normalization='none'): 123 | super(Upsampling, self).__init__() 124 | 125 | ops = [] 126 | ops.append(nn.Upsample(scale_factor=stride, mode='trilinear',align_corners=False)) 127 | ops.append(nn.Conv3d(n_filters_in, n_filters_out, kernel_size=3, padding=1)) 128 | if normalization == 'batchnorm': 129 | ops.append(nn.BatchNorm3d(n_filters_out)) 130 | elif normalization == 'groupnorm': 131 | ops.append(nn.GroupNorm(num_groups=16, num_channels=n_filters_out)) 132 | elif normalization == 'instancenorm': 133 | ops.append(nn.InstanceNorm3d(n_filters_out)) 134 | elif normalization != 'none': 135 | assert False 136 | ops.append(nn.ReLU(inplace=True)) 137 | 138 | self.conv = nn.Sequential(*ops) 139 | 140 | def forward(self, x): 141 | x = self.conv(x) 142 | return x 143 | 144 | 145 | class VNet(nn.Module): 146 | def __init__(self, n_channels=3, n_classes=2, n_filters=16, normalization='none', has_dropout=False, pyramid_has_dropout=False): 147 | super(VNet, self).__init__() 148 | self.has_dropout = has_dropout 149 | self.pyramid_has_dropout = pyramid_has_dropout 150 | 151 | self.block_one = ConvBlock(1, n_channels, n_filters, normalization=normalization) 152 | self.block_one_dw = DownsamplingConvBlock(n_filters, 2 * n_filters, normalization=normalization) 153 | 154 | self.block_two = ConvBlock(2, n_filters * 2, n_filters * 2, normalization=normalization) 155 | self.block_two_dw = DownsamplingConvBlock(n_filters * 2, n_filters * 4, normalization=normalization) 156 | 157 | self.block_three = ConvBlock(3, n_filters * 4, n_filters * 4, normalization=normalization) 158 | self.block_three_dw = DownsamplingConvBlock(n_filters * 4, n_filters * 8, normalization=normalization) 159 | 160 | self.block_four = ConvBlock(3, n_filters * 8, n_filters * 8, normalization=normalization) 161 | self.block_four_dw = DownsamplingConvBlock(n_filters * 8, n_filters * 16, normalization=normalization) 162 | 163 | self.block_five = ConvBlock(3, n_filters * 16, n_filters * 16, normalization=normalization) 164 | self.block_five_up = UpsamplingDeconvBlock(n_filters * 16, n_filters * 8, normalization=normalization) 165 | 166 | self.block_six = ConvBlock(3, n_filters * 8, n_filters * 8, normalization=normalization) 167 | self.block_six_up = UpsamplingDeconvBlock(n_filters * 8, n_filters * 4, normalization=normalization) 168 | 169 | self.block_seven = ConvBlock(3, n_filters * 4, n_filters * 4, normalization=normalization) 170 | self.block_seven_up = UpsamplingDeconvBlock(n_filters * 4, n_filters * 2, normalization=normalization) 171 | 172 | self.block_eight = ConvBlock(2, n_filters * 2, n_filters * 2, normalization=normalization) 173 | self.block_eight_up = UpsamplingDeconvBlock(n_filters * 2, n_filters, normalization=normalization) 174 | 175 | self.block_nine = ConvBlock(1, n_filters, n_filters, normalization=normalization) 176 | self.out_conv = nn.Conv3d(n_filters, n_classes, 1, padding=0) 177 | 178 | self.out_conv_dp3 = nn.Conv3d(n_filters * 8, n_classes, 1, padding=0) 179 | self.out_conv_dp2 = nn.Conv3d(n_filters * 4, n_classes, 1, padding=0) 180 | self.out_conv_dp1 = nn.Conv3d(n_filters * 2, n_classes, 1, padding=0) 181 | 182 | self.dropout = nn.Dropout3d(p=0.5, inplace=False) 183 | # self.__init_weight() 184 | 185 | def encoder(self, input): 186 | x1 = self.block_one(input) 187 | x1_dw = self.block_one_dw(x1) 188 | 189 | x2 = self.block_two(x1_dw) 190 | x2_dw = self.block_two_dw(x2) 191 | 192 | x3 = self.block_three(x2_dw) 193 | x3_dw = self.block_three_dw(x3) 194 | 195 | x4 = self.block_four(x3_dw) 196 | x4_dw = self.block_four_dw(x4) 197 | 198 | x5 = self.block_five(x4_dw) 199 | # x5 = F.dropout3d(x5, p=0.5, training=True) 200 | if self.has_dropout: 201 | x5 = self.dropout(x5) 202 | 203 | res = [x1, x2, x3, x4, x5] 204 | 205 | return res 206 | 207 | def decoder(self, features, shape): 208 | x1 = features[0] 209 | x2 = features[1] 210 | x3 = features[2] 211 | x4 = features[3] 212 | x5 = features[4] 213 | 214 | x5_up = self.block_five_up(x5) 215 | x5_up = x5_up + x4 216 | x6 = self.block_six(x5_up) 217 | if self.pyramid_has_dropout: 218 | dp3_out_seg = self.out_conv_dp3(self.dropout(x6)) 219 | else: 220 | dp3_out_seg = self.out_conv_dp3(x6) 221 | dp3_out_seg = F.interpolate(dp3_out_seg, shape) 222 | 223 | x6_up = self.block_six_up(x6) 224 | x6_up = x6_up + x3 225 | x7 = self.block_seven(x6_up) 226 | if self.pyramid_has_dropout: 227 | dp2_out_seg = self.out_conv_dp2(self.dropout(x7)) 228 | else: 229 | dp2_out_seg = self.out_conv_dp2(x7) 230 | dp2_out_seg = F.interpolate(dp2_out_seg, shape) 231 | 232 | x7_up = self.block_seven_up(x7) 233 | x7_up = x7_up + x2 234 | x8 = self.block_eight(x7_up) 235 | if self.pyramid_has_dropout: 236 | dp1_out_seg = self.out_conv_dp1(self.dropout(x8)) 237 | else: 238 | dp1_out_seg = self.out_conv_dp1(x8) 239 | dp1_out_seg = F.interpolate(dp1_out_seg, shape) 240 | 241 | x8_up = self.block_eight_up(x8) 242 | x8_up = x8_up + x1 243 | x9 = self.block_nine(x8_up) 244 | if self.has_dropout: 245 | x9 = self.dropout(x9) 246 | out = self.out_conv(x9) 247 | 248 | return out, dp1_out_seg, dp2_out_seg, dp3_out_seg 249 | 250 | 251 | def forward(self, input, turnoff_drop=False): 252 | if turnoff_drop: 253 | has_dropout = self.has_dropout 254 | self.has_dropout = False 255 | shape = input.shape[2:] 256 | features = self.encoder(input) 257 | dp1_out_seg, dp2_out_seg, dp3_out_seg, dp4_out_seg = self.decoder(features, shape) 258 | if turnoff_drop: 259 | self.has_dropout = has_dropout 260 | return dp1_out_seg, dp2_out_seg, dp3_out_seg, dp4_out_seg 261 | 262 | -------------------------------------------------------------------------------- /code/segment_anything_lora/predictor.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import numpy as np 8 | import torch 9 | 10 | from segment_anything.modeling import Sam 11 | 12 | from typing import Optional, Tuple 13 | 14 | from .utils.transforms import ResizeLongestSide 15 | 16 | 17 | class SamPredictor: 18 | def __init__( 19 | self, 20 | sam_model: Sam, 21 | ) -> None: 22 | """ 23 | Uses SAM to calculate the image embedding for an image, and then 24 | allow repeated, efficient mask prediction given prompts. 25 | 26 | Arguments: 27 | sam_model (Sam): The model to use for mask prediction. 28 | """ 29 | super().__init__() 30 | self.model = sam_model 31 | self.transform = ResizeLongestSide(sam_model.image_encoder.img_size) 32 | self.reset_image() 33 | 34 | def set_image( 35 | self, 36 | image: np.ndarray, 37 | image_format: str = "RGB", 38 | ) -> None: 39 | """ 40 | Calculates the image embeddings for the provided image, allowing 41 | masks to be predicted with the 'predict' method. 42 | 43 | Arguments: 44 | image (np.ndarray): The image for calculating masks. Expects an 45 | image in HWC uint8 format, with pixel values in [0, 255]. 46 | image_format (str): The color format of the image, in ['RGB', 'BGR']. 47 | """ 48 | assert image_format in [ 49 | "RGB", 50 | "BGR", 51 | ], f"image_format must be in ['RGB', 'BGR'], is {image_format}." 52 | if image_format != self.model.image_format: 53 | image = image[..., ::-1] 54 | 55 | # Transform the image to the form expected by the model 56 | input_image = self.transform.apply_image(image) 57 | input_image_torch = torch.as_tensor(input_image, device=self.device) 58 | input_image_torch = input_image_torch.permute(2, 0, 1).contiguous()[None, :, :, :] 59 | 60 | self.set_torch_image(input_image_torch, image.shape[:2]) 61 | 62 | @torch.no_grad() 63 | def set_torch_image( 64 | self, 65 | transformed_image: torch.Tensor, 66 | original_image_size: Tuple[int, ...], 67 | ) -> None: 68 | """ 69 | Calculates the image embeddings for the provided image, allowing 70 | masks to be predicted with the 'predict' method. Expects the input 71 | image to be already transformed to the format expected by the model. 72 | 73 | Arguments: 74 | transformed_image (torch.Tensor): The input image, with shape 75 | 1x3xHxW, which has been transformed with ResizeLongestSide. 76 | original_image_size (tuple(int, int)): The size of the image 77 | before transformation, in (H, W) format. 78 | """ 79 | assert ( 80 | len(transformed_image.shape) == 4 81 | and transformed_image.shape[1] == 3 82 | and max(*transformed_image.shape[2:]) == self.model.image_encoder.img_size 83 | ), f"set_torch_image input must be BCHW with long side {self.model.image_encoder.img_size}." 84 | self.reset_image() 85 | 86 | self.original_size = original_image_size 87 | self.input_size = tuple(transformed_image.shape[-2:]) 88 | input_image = self.model.preprocess(transformed_image) 89 | self.features = self.model.image_encoder(input_image) 90 | self.is_image_set = True 91 | 92 | def predict( 93 | self, 94 | point_coords: Optional[np.ndarray] = None, 95 | point_labels: Optional[np.ndarray] = None, 96 | box: Optional[np.ndarray] = None, 97 | mask_input: Optional[np.ndarray] = None, 98 | multimask_output: bool = True, 99 | return_logits: bool = False, 100 | ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: 101 | """ 102 | Predict masks for the given input prompts, using the currently set image. 103 | 104 | Arguments: 105 | point_coords (np.ndarray or None): A Nx2 array of point prompts to the 106 | model. Each point is in (X,Y) in pixels. 107 | point_labels (np.ndarray or None): A length N array of labels for the 108 | point prompts. 1 indicates a foreground point and 0 indicates a 109 | background point. 110 | box (np.ndarray or None): A length 4 array given a box prompt to the 111 | model, in XYXY format. 112 | mask_input (np.ndarray): A low resolution mask input to the model, typically 113 | coming from a previous prediction iteration. Has form 1xHxW, where 114 | for SAM, H=W=256. 115 | multimask_output (bool): If true, the model will return three masks. 116 | For ambiguous input prompts (such as a single click), this will often 117 | produce better masks than a single prediction. If only a single 118 | mask is needed, the model's predicted quality score can be used 119 | to select the best mask. For non-ambiguous prompts, such as multiple 120 | input prompts, multimask_output=False can give better results. 121 | return_logits (bool): If true, returns un-thresholded masks logits 122 | instead of a binary mask. 123 | 124 | Returns: 125 | (np.ndarray): The output masks in CxHxW format, where C is the 126 | number of masks, and (H, W) is the original image size. 127 | (np.ndarray): An array of length C containing the model's 128 | predictions for the quality of each mask. 129 | (np.ndarray): An array of shape CxHxW, where C is the number 130 | of masks and H=W=256. These low resolution logits can be passed to 131 | a subsequent iteration as mask input. 132 | """ 133 | if not self.is_image_set: 134 | raise RuntimeError("An image must be set with .set_image(...) before mask prediction.") 135 | 136 | # Transform input prompts 137 | coords_torch, labels_torch, box_torch, mask_input_torch = None, None, None, None 138 | if point_coords is not None: 139 | assert ( 140 | point_labels is not None 141 | ), "point_labels must be supplied if point_coords is supplied." 142 | point_coords = self.transform.apply_coords(point_coords, self.original_size) 143 | coords_torch = torch.as_tensor(point_coords, dtype=torch.float, device=self.device) 144 | labels_torch = torch.as_tensor(point_labels, dtype=torch.int, device=self.device) 145 | coords_torch, labels_torch = coords_torch[None, :, :], labels_torch[None, :] 146 | if box is not None: 147 | box = self.transform.apply_boxes(box, self.original_size) 148 | box_torch = torch.as_tensor(box, dtype=torch.float, device=self.device) 149 | box_torch = box_torch[None, :] 150 | if mask_input is not None: 151 | mask_input_torch = torch.as_tensor(mask_input, dtype=torch.float, device=self.device) 152 | mask_input_torch = mask_input_torch[None, :, :, :] 153 | 154 | masks, iou_predictions, low_res_masks = self.predict_torch( 155 | coords_torch, 156 | labels_torch, 157 | box_torch, 158 | mask_input_torch, 159 | multimask_output, 160 | return_logits=return_logits, 161 | ) 162 | 163 | masks = masks[0].detach().cpu().numpy() 164 | iou_predictions = iou_predictions[0].detach().cpu().numpy() 165 | low_res_masks = low_res_masks[0].detach().cpu().numpy() 166 | return masks, iou_predictions, low_res_masks 167 | 168 | @torch.no_grad() 169 | def predict_torch( 170 | self, 171 | point_coords: Optional[torch.Tensor], 172 | point_labels: Optional[torch.Tensor], 173 | boxes: Optional[torch.Tensor] = None, 174 | mask_input: Optional[torch.Tensor] = None, 175 | multimask_output: bool = True, 176 | return_logits: bool = False, 177 | ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: 178 | """ 179 | Predict masks for the given input prompts, using the currently set image. 180 | Input prompts are batched torch tensors and are expected to already be 181 | transformed to the input frame using ResizeLongestSide. 182 | 183 | Arguments: 184 | point_coords (torch.Tensor or None): A BxNx2 array of point prompts to the 185 | model. Each point is in (X,Y) in pixels. 186 | point_labels (torch.Tensor or None): A BxN array of labels for the 187 | point prompts. 1 indicates a foreground point and 0 indicates a 188 | background point. 189 | box (np.ndarray or None): A Bx4 array given a box prompt to the 190 | model, in XYXY format. 191 | mask_input (np.ndarray): A low resolution mask input to the model, typically 192 | coming from a previous prediction iteration. Has form Bx1xHxW, where 193 | for SAM, H=W=256. Masks returned by a previous iteration of the 194 | predict method do not need further transformation. 195 | multimask_output (bool): If true, the model will return three masks. 196 | For ambiguous input prompts (such as a single click), this will often 197 | produce better masks than a single prediction. If only a single 198 | mask is needed, the model's predicted quality score can be used 199 | to select the best mask. For non-ambiguous prompts, such as multiple 200 | input prompts, multimask_output=False can give better results. 201 | return_logits (bool): If true, returns un-thresholded masks logits 202 | instead of a binary mask. 203 | 204 | Returns: 205 | (torch.Tensor): The output masks in BxCxHxW format, where C is the 206 | number of masks, and (H, W) is the original image size. 207 | (torch.Tensor): An array of shape BxC containing the model's 208 | predictions for the quality of each mask. 209 | (torch.Tensor): An array of shape BxCxHxW, where C is the number 210 | of masks and H=W=256. These low res logits can be passed to 211 | a subsequent iteration as mask input. 212 | """ 213 | if not self.is_image_set: 214 | raise RuntimeError("An image must be set with .set_image(...) before mask prediction.") 215 | 216 | if point_coords is not None: 217 | points = (point_coords, point_labels) 218 | else: 219 | points = None 220 | 221 | # Embed prompts 222 | sparse_embeddings, dense_embeddings = self.model.prompt_encoder( 223 | points=points, 224 | boxes=boxes, 225 | masks=mask_input, 226 | ) 227 | 228 | # Predict masks 229 | low_res_masks, iou_predictions = self.model.mask_decoder( 230 | image_embeddings=self.features, 231 | image_pe=self.model.prompt_encoder.get_dense_pe(), 232 | sparse_prompt_embeddings=sparse_embeddings, 233 | dense_prompt_embeddings=dense_embeddings, 234 | multimask_output=multimask_output, 235 | ) 236 | 237 | # Upscale the masks to the original image resolution 238 | masks = self.model.postprocess_masks(low_res_masks, self.input_size, self.original_size) 239 | 240 | if not return_logits: 241 | masks = masks > self.model.mask_threshold 242 | 243 | return masks, iou_predictions, low_res_masks 244 | 245 | def get_image_embedding(self) -> torch.Tensor: 246 | """ 247 | Returns the image embeddings for the currently set image, with 248 | shape 1xCxHxW, where C is the embedding dimension and (H,W) are 249 | the embedding spatial dimension of SAM (typically C=256, H=W=64). 250 | """ 251 | if not self.is_image_set: 252 | raise RuntimeError( 253 | "An image must be set with .set_image(...) to generate an embedding." 254 | ) 255 | assert self.features is not None, "Features must exist if an image has been set." 256 | return self.features 257 | 258 | @property 259 | def device(self) -> torch.device: 260 | return self.model.device 261 | 262 | def reset_image(self) -> None: 263 | """Resets the currently set image.""" 264 | self.is_image_set = False 265 | self.features = None 266 | self.orig_h = None 267 | self.orig_w = None 268 | self.input_h = None 269 | self.input_w = None 270 | -------------------------------------------------------------------------------- /code/train_finetuning.py: -------------------------------------------------------------------------------- 1 | import os 2 | import pdb 3 | import sys 4 | import random 5 | import shutil 6 | import argparse 7 | import logging 8 | import numpy as np 9 | import nibabel as nib 10 | from tqdm import tqdm 11 | from scipy import ndimage 12 | from importlib import import_module 13 | from tensorboardX import SummaryWriter 14 | 15 | parser = argparse.ArgumentParser() 16 | parser.add_argument('--root_path', type=str, default='../data/LA/processed_h5_rdm_4/', help='Name of Experiment') 17 | parser.add_argument('--exp', type=str, default='name', help='model_name') 18 | parser.add_argument('--dataset', type=str, default='la', help='dataset to use') 19 | parser.add_argument('--label_num', type=int, default=16, help='number of labeled data') 20 | 21 | parser.add_argument('--pretrain_model', type=str, default='vit_b', help='vit to select') 22 | parser.add_argument('--patch_size', type=int, default=128, help='shape of data') 23 | parser.add_argument('--input_size', type=int, default=1024, help='shape of data') 24 | parser.add_argument('--num_classes', type=int, default=1, help='number of class') 25 | parser.add_argument('--save_img', type=int, default=250, help='img saving iterations') 26 | # load 27 | parser.add_argument('--load', action="store_true", help='load net') 28 | parser.add_argument('--load_iter', type=int, default=0, help='load iter') 29 | 30 | parser.add_argument('--save_iter', type=int, default=1000, help='maximum epoch number to train') 31 | parser.add_argument('--max_iterations', type=int, default=6000, help='maximum epoch number to train') 32 | parser.add_argument('--batch_size', type=int, default=2, help='batch_size per gpu') 33 | 34 | parser.add_argument('--nt', type=int, default=2, help='nonlinear_transformation') 35 | parser.add_argument('--nonlinear_rate', type=float, default=0.5, help='nonlinear_rate') 36 | parser.add_argument('--rdmrotflip', action="store_true", help='rdmrotflip') 37 | 38 | parser.add_argument("--lr_sam", type=float, default=0.001, help="sam learning rate") 39 | parser.add_argument('--warmup', action='store_true', help='If activated, warp up the learning from a lower lr to the base_lr') 40 | parser.add_argument('--warmup_period', type=int, default=250, help='Warp up iterations, only valid whrn warmup is activated') 41 | 42 | parser.add_argument('--deterministic', type=int, default=1, help='whether use deterministic training') 43 | parser.add_argument('--seed', type=int, default=1337, help='random seed') 44 | parser.add_argument('--gpu', type=str, default='0', help='GPU to use') 45 | parser.add_argument('--rank', type=int, default=4, help='Rank for LoRA adaptation') 46 | parser.add_argument('--module', type=str, default='sam_lora_image_encoder') 47 | args = parser.parse_args() 48 | 49 | root = "../" 50 | 51 | train_data_path = args.root_path 52 | snapshot_path = root + "model_" + args.dataset + "/" + args.exp + "/" 53 | 54 | os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu 55 | batch_size = args.batch_size 56 | n_gpu = len(args.gpu.split(',')) 57 | print(batch_size) 58 | max_iterations, input_size, patch_size = args.max_iterations, args.input_size, args.patch_size 59 | num_classes = args.num_classes 60 | lr_sam = args.lr_sam 61 | 62 | import torch 63 | import torch.optim as optim 64 | import torch.nn.functional as F 65 | import torch.backends.cudnn as cudnn 66 | from torchvision import transforms 67 | from torch.utils.data import DataLoader 68 | from torch.nn.modules.loss import CrossEntropyLoss 69 | 70 | from sam_lora_image_encoder import LoRA_Sam 71 | from segment_anything_lora import sam_model_registry 72 | from dataloaders.dataset import * 73 | from utils import ramps, losses 74 | from utils.util import * 75 | 76 | if args.deterministic: 77 | cudnn.benchmark = False 78 | cudnn.deterministic = True 79 | random.seed(args.seed) 80 | np.random.seed(args.seed) 81 | torch.manual_seed(args.seed) 82 | torch.cuda.manual_seed(args.seed) 83 | 84 | def worker_init_fn(worker_id): 85 | random.seed(args.seed + worker_id) 86 | 87 | 88 | if __name__ == "__main__": 89 | ## make logger file 90 | if not os.path.exists(snapshot_path): 91 | os.makedirs(snapshot_path) 92 | os.makedirs(snapshot_path + 'saveimg') 93 | 94 | logging.basicConfig(filename=snapshot_path + "/log.txt", level=logging.INFO, 95 | format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S') 96 | logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) 97 | logging.info(str(args)) 98 | 99 | if args.dataset == 'la': 100 | db_train = LAHeart(base_dir=train_data_path, 101 | num=args.label_num, 102 | transform=transforms.Compose([ 103 | RandomRotFlip(), 104 | RandomCrop((112, 112, 80)), 105 | ToTensor(), 106 | ])) 107 | 108 | elif args.dataset == 'btcv': 109 | db_train = BTCV(base_dir=train_data_path, 110 | num=args.label_num, 111 | transform=transforms.Compose([ 112 | RandomCrop((patch_size, patch_size, patch_size)), 113 | ToTensor(), 114 | ])) 115 | 116 | elif args.dataset == 'mact': 117 | db_train = MACT(base_dir=train_data_path, 118 | num=args.label_num, 119 | transform=transforms.Compose([ 120 | RandomCrop((patch_size, patch_size, patch_size)), 121 | ToTensor(), 122 | ])) 123 | 124 | elif args.dataset == 'brats': 125 | db_train = BraTS19(base_dir=train_data_path, 126 | num=args.label_num, 127 | transform=transforms.Compose([ 128 | RandomRotFlip(), 129 | RandomCrop((patch_size, patch_size, patch_size)), 130 | ToTensor(), 131 | ])) 132 | 133 | multimask_output = True if num_classes > 2 else False 134 | trainloader = DataLoader(db_train, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True, worker_init_fn=worker_init_fn) 135 | 136 | if args.pretrain_model == "vit_b": 137 | model_sam, img_embedding_size = sam_model_registry["vit_b"](image_size=args.input_size, num_classes=num_classes, checkpoint='pre_weight/sam_vit_b_01ec64.pth', pixel_mean=[0, 0, 0], pixel_std=[1, 1, 1]) 138 | pkg = import_module(args.module) 139 | model_sam = pkg.LoRA_Sam(model_sam, args.rank).cuda() 140 | if args.load: 141 | sam_checkpoint = snapshot_path + "/sam_iter_" + str(args.load_iter) + ".pth" 142 | model_sam.load_lora_parameters(sam_checkpoint) 143 | print("init weight from {}".format(sam_checkpoint)) 144 | if args.warmup: 145 | base_lr_sam = lr_sam / args.warmup_period 146 | else: 147 | base_lr_sam = lr_sam 148 | optimizer_sam = torch.optim.AdamW(filter(lambda p: p.requires_grad, model_sam.parameters()), lr=base_lr_sam, betas=(0.9, 0.999), weight_decay=0.1) 149 | model_sam.train() 150 | 151 | # Set losses 152 | ce_loss = CrossEntropyLoss(ignore_index=255) 153 | dice_loss = losses.DiceLoss(num_classes+1) 154 | 155 | writer = SummaryWriter(snapshot_path + '/log') 156 | logging.info("{} itertations per epoch".format(len(trainloader))) 157 | 158 | iter_num = args.load_iter 159 | max_epoch = (max_iterations - args.load_iter) // len(trainloader) + 1 160 | kl_distance = torch.nn.KLDivLoss(reduction='none') 161 | lr_ = base_lr_sam 162 | 163 | for epoch_num in range(max_epoch): 164 | for i_batch, sampled_batch in enumerate(trainloader): 165 | volume_batch, label_batch = sampled_batch['image'], sampled_batch['label'] # [2, 1, 128, 128, 64], [2, 128, 128, 64] 166 | 167 | ### Train SAM Module 168 | sam_volume_batch = volume_batch.cpu().detach().numpy() 169 | sam_label_batch = label_batch.cpu().detach().numpy() 170 | 171 | ## labeled data 172 | image = sam_volume_batch # [B, 1, 128, 128, 64] 173 | label = sam_label_batch # [B, 128, 128, 64] 174 | 175 | image_inshape, label_inshape, n_row, n_col, pw, ph, ps, pww, phh, s_l, s_w = Spread_bs_aug(image, label, input_size, args.nt, args.nonlinear_rate) # (B, 1024, 1024) 176 | 177 | if args.rdmrotflip: # 2d RandomRotFlip 178 | k = np.array([np.random.randint(0, 4) for _ in range(image_inshape.shape[0])]) 179 | axis = np.array([np.random.randint(0, 2) for _ in range(image_inshape.shape[0])]) 180 | for i in range(image_inshape.shape[0]): 181 | image_inshape[i] = RandomRotFlip_2d(image_inshape[i], k[i], axis[i]) 182 | label_inshape[i] = RandomRotFlip_2d(label_inshape[i], k[i], axis[i]) 183 | 184 | volume_batch_inshape, label_batch_inshape = ToTensor_sam_bs(image_inshape, label_inshape) # [B, 3, 1024, 1024], [B, 1024, 1024] 185 | 186 | outputs = model_sam(volume_batch_inshape, multimask_output, input_size) 187 | output_masks, low_res_masks, iou_predictions = outputs['masks'], outputs['low_res_logits'], outputs['iou_predictions'] # [1, C=2/14, 1024, 1024], [1, 2, 256, 256], [1, 2] 188 | output_soft = F.softmax(output_masks, dim=1) # [1, 2, 1024, 1024] 189 | 190 | ## calculate the loss - labeled data 191 | loss_ce_sam = ce_loss(output_masks, label_batch_inshape) 192 | loss_dice_sam = dice_loss(output_soft, label_batch_inshape.unsqueeze(1)) 193 | loss_sam = 0.5 * (loss_ce_sam + loss_dice_sam) 194 | optimizer_sam.zero_grad() 195 | loss_sam.backward() 196 | optimizer_sam.step() 197 | 198 | if iter_num % args.save_img == 0: 199 | prediction_inshape = torch.argmax(output_soft, dim=1) # (B, 1024, 1024) 200 | prediction_inshape = prediction_inshape.cpu().detach().numpy() # (B, 1024, 1024) 201 | 202 | nib.save(nib.Nifti1Image(image_inshape.astype(np.float32), np.eye(4)), snapshot_path + '/saveimg' + '/inshape_img_' + str(iter_num) + '.nii.gz') 203 | nib.save(nib.Nifti1Image(label_inshape.astype(np.float32), np.eye(4)), snapshot_path + '/saveimg' + '/inshape_gt_' + str(iter_num) + '.nii.gz') 204 | nib.save(nib.Nifti1Image(prediction_inshape.astype(np.float32), np.eye(4)), snapshot_path + '/saveimg' + '/inshape_pred_' + str(iter_num) + '_u.nii.gz') 205 | 206 | if args.rdmrotflip: 207 | for i in range(image_inshape.shape[0]): 208 | prediction_inshape[i] = RandomFlipRot_2d(prediction_inshape[i], k[i]*(-1), axis[i]*(-1)) 209 | 210 | pred_single = np.zeros((batch_size, image.shape[-1] + ps * 2, image.shape[-2] + pw * 2, image.shape[-3] + ph * 2)) # [B, 80, 112, 112] (1, 112, 114, 86) 211 | if pww < 0 or phh < 0: 212 | pww_r, phh_r = pww, phh 213 | prediction_inshape = np.pad(prediction_inshape, [(0, 0), (abs(pww), abs(pww)), (abs(phh), abs(phh))], mode='constant', constant_values=255) 214 | pww, phh = 0, 0 215 | for row in range(n_row): 216 | for col in range(n_col): 217 | if row * n_col + col < pred_single.shape[1]: 218 | pred_single[:, row * n_col + col] = prediction_inshape[:, pww + row * s_l:pww + (row + 1) * s_l, phh + col * s_w:phh + (col + 1) * s_w] 219 | pred_single = pred_single[:, ps:pred_single.shape[-3] - ps, pw:pred_single.shape[-2] - pw, ph:pred_single.shape[-1] - ph] # (B, 64, 128, 128) 220 | pred_single = np.swapaxes(pred_single, -3, -1) # (B, 128, 128, 64) 221 | prediction_sam = pred_single 222 | 223 | nib.save(nib.Nifti1Image(image[0,0].astype(np.float32), np.eye(4)), snapshot_path + '/saveimg' + '/img_' + str(iter_num) + '.nii.gz') 224 | nib.save(nib.Nifti1Image(label[0].astype(np.float32), np.eye(4)), snapshot_path + '/saveimg' + '/gt_' + str(iter_num) + '.nii.gz') 225 | nib.save(nib.Nifti1Image(prediction_sam[0].astype(np.float32), np.eye(4)), snapshot_path + '/saveimg' + '/pred_sam_' + str(iter_num) + '_u.nii.gz') 226 | 227 | logging.info('iter %d : sam loss : %f, ce loss : %f, dice loss : %f' % (iter_num, loss_sam, loss_ce_sam, loss_dice_sam.item())) 228 | 229 | iter_num = iter_num + 1 230 | writer.add_scalar('lr', lr_, iter_num) 231 | writer.add_scalar('loss/loss_sam', loss_sam, iter_num) 232 | 233 | ## change lr 234 | if args.warmup and iter_num < args.warmup_period: 235 | lr_ = lr_sam * ((iter_num + 1) / args.warmup_period) 236 | else: 237 | if args.warmup: 238 | shift_iter = iter_num - args.warmup_period # assert shift_iter >= 0, f'Shift iter is {shift_iter}, smaller than zero' 239 | else: 240 | shift_iter = iter_num 241 | lr_ = lr_sam * (1.0 - shift_iter / max_iterations) ** 0.9 # learning rate adjustment depends on the max iterations 242 | for param_group in optimizer_sam.param_groups: 243 | param_group['lr'] = lr_ 244 | 245 | if iter_num % args.save_iter == 0: 246 | save_mode_path_sam = os.path.join(snapshot_path, 'sam_iter_' + str(iter_num) + '.pth') 247 | model_sam.save_lora_parameters(save_mode_path_sam) 248 | logging.info("save model to {}".format(save_mode_path_sam)) 249 | 250 | if iter_num >= max_iterations: 251 | break 252 | if iter_num >= max_iterations: 253 | break 254 | save_mode_path_sam = os.path.join(snapshot_path, 'sam_iter_' + str(iter_num) + '.pth') 255 | model_sam.save_lora_parameters(save_mode_path_sam) 256 | logging.info("save model to {}".format(save_mode_path_sam)) 257 | writer.close() 258 | 259 | 260 | -------------------------------------------------------------------------------- /code/segment_anything_lora/utils/amg.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import numpy as np 8 | import torch 9 | 10 | import math 11 | from copy import deepcopy 12 | from itertools import product 13 | from typing import Any, Dict, Generator, ItemsView, List, Tuple 14 | 15 | 16 | class MaskData: 17 | """ 18 | A structure for storing masks and their related data in batched format. 19 | Implements basic filtering and concatenation. 20 | """ 21 | 22 | def __init__(self, **kwargs) -> None: 23 | for v in kwargs.values(): 24 | assert isinstance( 25 | v, (list, np.ndarray, torch.Tensor) 26 | ), "MaskData only supports list, numpy arrays, and torch tensors." 27 | self._stats = dict(**kwargs) 28 | 29 | def __setitem__(self, key: str, item: Any) -> None: 30 | assert isinstance( 31 | item, (list, np.ndarray, torch.Tensor) 32 | ), "MaskData only supports list, numpy arrays, and torch tensors." 33 | self._stats[key] = item 34 | 35 | def __delitem__(self, key: str) -> None: 36 | del self._stats[key] 37 | 38 | def __getitem__(self, key: str) -> Any: 39 | return self._stats[key] 40 | 41 | def items(self) -> ItemsView[str, Any]: 42 | return self._stats.items() 43 | 44 | def filter(self, keep: torch.Tensor) -> None: 45 | for k, v in self._stats.items(): 46 | if v is None: 47 | self._stats[k] = None 48 | elif isinstance(v, torch.Tensor): 49 | self._stats[k] = v[torch.as_tensor(keep, device=v.device)] 50 | elif isinstance(v, np.ndarray): 51 | self._stats[k] = v[keep.detach().cpu().numpy()] 52 | elif isinstance(v, list) and keep.dtype == torch.bool: 53 | self._stats[k] = [a for i, a in enumerate(v) if keep[i]] 54 | elif isinstance(v, list): 55 | self._stats[k] = [v[i] for i in keep] 56 | else: 57 | raise TypeError(f"MaskData key {k} has an unsupported type {type(v)}.") 58 | 59 | def cat(self, new_stats: "MaskData") -> None: 60 | for k, v in new_stats.items(): 61 | if k not in self._stats or self._stats[k] is None: 62 | self._stats[k] = deepcopy(v) 63 | elif isinstance(v, torch.Tensor): 64 | self._stats[k] = torch.cat([self._stats[k], v], dim=0) 65 | elif isinstance(v, np.ndarray): 66 | self._stats[k] = np.concatenate([self._stats[k], v], axis=0) 67 | elif isinstance(v, list): 68 | self._stats[k] = self._stats[k] + deepcopy(v) 69 | else: 70 | raise TypeError(f"MaskData key {k} has an unsupported type {type(v)}.") 71 | 72 | def to_numpy(self) -> None: 73 | for k, v in self._stats.items(): 74 | if isinstance(v, torch.Tensor): 75 | self._stats[k] = v.detach().cpu().numpy() 76 | 77 | 78 | def is_box_near_crop_edge( 79 | boxes: torch.Tensor, crop_box: List[int], orig_box: List[int], atol: float = 20.0 80 | ) -> torch.Tensor: 81 | """Filter masks at the edge of a crop, but not at the edge of the original image.""" 82 | crop_box_torch = torch.as_tensor(crop_box, dtype=torch.float, device=boxes.device) 83 | orig_box_torch = torch.as_tensor(orig_box, dtype=torch.float, device=boxes.device) 84 | boxes = uncrop_boxes_xyxy(boxes, crop_box).float() 85 | near_crop_edge = torch.isclose(boxes, crop_box_torch[None, :], atol=atol, rtol=0) 86 | near_image_edge = torch.isclose(boxes, orig_box_torch[None, :], atol=atol, rtol=0) 87 | near_crop_edge = torch.logical_and(near_crop_edge, ~near_image_edge) 88 | return torch.any(near_crop_edge, dim=1) 89 | 90 | 91 | def box_xyxy_to_xywh(box_xyxy: torch.Tensor) -> torch.Tensor: 92 | box_xywh = deepcopy(box_xyxy) 93 | box_xywh[2] = box_xywh[2] - box_xywh[0] 94 | box_xywh[3] = box_xywh[3] - box_xywh[1] 95 | return box_xywh 96 | 97 | 98 | def batch_iterator(batch_size: int, *args) -> Generator[List[Any], None, None]: 99 | assert len(args) > 0 and all( 100 | len(a) == len(args[0]) for a in args 101 | ), "Batched iteration must have inputs of all the same size." 102 | n_batches = len(args[0]) // batch_size + int(len(args[0]) % batch_size != 0) 103 | for b in range(n_batches): 104 | yield [arg[b * batch_size : (b + 1) * batch_size] for arg in args] 105 | 106 | 107 | def mask_to_rle_pytorch(tensor: torch.Tensor) -> List[Dict[str, Any]]: 108 | """ 109 | Encodes masks to an uncompressed RLE, in the format expected by 110 | pycoco tools. 111 | """ 112 | # Put in fortran order and flatten h,w 113 | b, h, w = tensor.shape 114 | tensor = tensor.permute(0, 2, 1).flatten(1) 115 | 116 | # Compute change indices 117 | diff = tensor[:, 1:] ^ tensor[:, :-1] 118 | change_indices = diff.nonzero() 119 | 120 | # Encode run length 121 | out = [] 122 | for i in range(b): 123 | cur_idxs = change_indices[change_indices[:, 0] == i, 1] 124 | cur_idxs = torch.cat( 125 | [ 126 | torch.tensor([0], dtype=cur_idxs.dtype, device=cur_idxs.device), 127 | cur_idxs + 1, 128 | torch.tensor([h * w], dtype=cur_idxs.dtype, device=cur_idxs.device), 129 | ] 130 | ) 131 | btw_idxs = cur_idxs[1:] - cur_idxs[:-1] 132 | counts = [] if tensor[i, 0] == 0 else [0] 133 | counts.extend(btw_idxs.detach().cpu().tolist()) 134 | out.append({"size": [h, w], "counts": counts}) 135 | return out 136 | 137 | 138 | def rle_to_mask(rle: Dict[str, Any]) -> np.ndarray: 139 | """Compute a binary mask from an uncompressed RLE.""" 140 | h, w = rle["size"] 141 | mask = np.empty(h * w, dtype=bool) 142 | idx = 0 143 | parity = False 144 | for count in rle["counts"]: 145 | mask[idx : idx + count] = parity 146 | idx += count 147 | parity ^= True 148 | mask = mask.reshape(w, h) 149 | return mask.transpose() # Put in C order 150 | 151 | 152 | def area_from_rle(rle: Dict[str, Any]) -> int: 153 | return sum(rle["counts"][1::2]) 154 | 155 | 156 | def calculate_stability_score( 157 | masks: torch.Tensor, mask_threshold: float, threshold_offset: float 158 | ) -> torch.Tensor: 159 | """ 160 | Computes the stability score for a batch of masks. The stability 161 | score is the IoU between the binary masks obtained by thresholding 162 | the predicted mask logits at high and low values. 163 | """ 164 | # One mask is always contained inside the other. 165 | # Save memory by preventing unnecesary cast to torch.int64 166 | intersections = ( 167 | (masks > (mask_threshold + threshold_offset)) 168 | .sum(-1, dtype=torch.int16) 169 | .sum(-1, dtype=torch.int32) 170 | ) 171 | unions = ( 172 | (masks > (mask_threshold - threshold_offset)) 173 | .sum(-1, dtype=torch.int16) 174 | .sum(-1, dtype=torch.int32) 175 | ) 176 | return intersections / unions 177 | 178 | 179 | def build_point_grid(n_per_side: int) -> np.ndarray: 180 | """Generates a 2D grid of points evenly spaced in [0,1]x[0,1].""" 181 | offset = 1 / (2 * n_per_side) 182 | points_one_side = np.linspace(offset, 1 - offset, n_per_side) 183 | points_x = np.tile(points_one_side[None, :], (n_per_side, 1)) 184 | points_y = np.tile(points_one_side[:, None], (1, n_per_side)) 185 | points = np.stack([points_x, points_y], axis=-1).reshape(-1, 2) 186 | return points 187 | 188 | 189 | def build_all_layer_point_grids( 190 | n_per_side: int, n_layers: int, scale_per_layer: int 191 | ) -> List[np.ndarray]: 192 | """Generates point grids for all crop layers.""" 193 | points_by_layer = [] 194 | for i in range(n_layers + 1): 195 | n_points = int(n_per_side / (scale_per_layer**i)) 196 | points_by_layer.append(build_point_grid(n_points)) 197 | return points_by_layer 198 | 199 | 200 | def generate_crop_boxes( 201 | im_size: Tuple[int, ...], n_layers: int, overlap_ratio: float 202 | ) -> Tuple[List[List[int]], List[int]]: 203 | """ 204 | Generates a list of crop boxes of different sizes. Each layer 205 | has (2**i)**2 boxes for the ith layer. 206 | """ 207 | crop_boxes, layer_idxs = [], [] 208 | im_h, im_w = im_size 209 | short_side = min(im_h, im_w) 210 | 211 | # Original image 212 | crop_boxes.append([0, 0, im_w, im_h]) 213 | layer_idxs.append(0) 214 | 215 | def crop_len(orig_len, n_crops, overlap): 216 | return int(math.ceil((overlap * (n_crops - 1) + orig_len) / n_crops)) 217 | 218 | for i_layer in range(n_layers): 219 | n_crops_per_side = 2 ** (i_layer + 1) 220 | overlap = int(overlap_ratio * short_side * (2 / n_crops_per_side)) 221 | 222 | crop_w = crop_len(im_w, n_crops_per_side, overlap) 223 | crop_h = crop_len(im_h, n_crops_per_side, overlap) 224 | 225 | crop_box_x0 = [int((crop_w - overlap) * i) for i in range(n_crops_per_side)] 226 | crop_box_y0 = [int((crop_h - overlap) * i) for i in range(n_crops_per_side)] 227 | 228 | # Crops in XYWH format 229 | for x0, y0 in product(crop_box_x0, crop_box_y0): 230 | box = [x0, y0, min(x0 + crop_w, im_w), min(y0 + crop_h, im_h)] 231 | crop_boxes.append(box) 232 | layer_idxs.append(i_layer + 1) 233 | 234 | return crop_boxes, layer_idxs 235 | 236 | 237 | def uncrop_boxes_xyxy(boxes: torch.Tensor, crop_box: List[int]) -> torch.Tensor: 238 | x0, y0, _, _ = crop_box 239 | offset = torch.tensor([[x0, y0, x0, y0]], device=boxes.device) 240 | # Check if boxes has a channel dimension 241 | if len(boxes.shape) == 3: 242 | offset = offset.unsqueeze(1) 243 | return boxes + offset 244 | 245 | 246 | def uncrop_points(points: torch.Tensor, crop_box: List[int]) -> torch.Tensor: 247 | x0, y0, _, _ = crop_box 248 | offset = torch.tensor([[x0, y0]], device=points.device) 249 | # Check if points has a channel dimension 250 | if len(points.shape) == 3: 251 | offset = offset.unsqueeze(1) 252 | return points + offset 253 | 254 | 255 | def uncrop_masks( 256 | masks: torch.Tensor, crop_box: List[int], orig_h: int, orig_w: int 257 | ) -> torch.Tensor: 258 | x0, y0, x1, y1 = crop_box 259 | if x0 == 0 and y0 == 0 and x1 == orig_w and y1 == orig_h: 260 | return masks 261 | # Coordinate transform masks 262 | pad_x, pad_y = orig_w - (x1 - x0), orig_h - (y1 - y0) 263 | pad = (x0, pad_x - x0, y0, pad_y - y0) 264 | return torch.nn.functional.pad(masks, pad, value=0) 265 | 266 | 267 | def remove_small_regions( 268 | mask: np.ndarray, area_thresh: float, mode: str 269 | ) -> Tuple[np.ndarray, bool]: 270 | """ 271 | Removes small disconnected regions and holes in a mask. Returns the 272 | mask and an indicator of if the mask has been modified. 273 | """ 274 | import cv2 # type: ignore 275 | 276 | assert mode in ["holes", "islands"] 277 | correct_holes = mode == "holes" 278 | working_mask = (correct_holes ^ mask).astype(np.uint8) 279 | n_labels, regions, stats, _ = cv2.connectedComponentsWithStats(working_mask, 8) 280 | sizes = stats[:, -1][1:] # Row 0 is background label 281 | small_regions = [i + 1 for i, s in enumerate(sizes) if s < area_thresh] 282 | if len(small_regions) == 0: 283 | return mask, False 284 | fill_labels = [0] + small_regions 285 | if not correct_holes: 286 | fill_labels = [i for i in range(n_labels) if i not in fill_labels] 287 | # If every region is below threshold, keep largest 288 | if len(fill_labels) == 0: 289 | fill_labels = [int(np.argmax(sizes)) + 1] 290 | mask = np.isin(regions, fill_labels) 291 | return mask, True 292 | 293 | 294 | def coco_encode_rle(uncompressed_rle: Dict[str, Any]) -> Dict[str, Any]: 295 | from pycocotools import mask as mask_utils # type: ignore 296 | 297 | h, w = uncompressed_rle["size"] 298 | rle = mask_utils.frPyObjects(uncompressed_rle, h, w) 299 | rle["counts"] = rle["counts"].decode("utf-8") # Necessary to serialize with json 300 | return rle 301 | 302 | 303 | def batched_mask_to_box(masks: torch.Tensor) -> torch.Tensor: 304 | """ 305 | Calculates boxes in XYXY format around masks. Return [0,0,0,0] for 306 | an empty mask. For input shape C1xC2x...xHxW, the output shape is C1xC2x...x4. 307 | """ 308 | # torch.max below raises an error on empty inputs, just skip in this case 309 | if torch.numel(masks) == 0: 310 | return torch.zeros(*masks.shape[:-2], 4, device=masks.device) 311 | 312 | # Normalize shape to CxHxW 313 | shape = masks.shape 314 | h, w = shape[-2:] 315 | if len(shape) > 2: 316 | masks = masks.flatten(0, -3) 317 | else: 318 | masks = masks.unsqueeze(0) 319 | 320 | # Get top and bottom edges 321 | in_height, _ = torch.max(masks, dim=-1) 322 | in_height_coords = in_height * torch.arange(h, device=in_height.device)[None, :] 323 | bottom_edges, _ = torch.max(in_height_coords, dim=-1) 324 | in_height_coords = in_height_coords + h * (~in_height) 325 | top_edges, _ = torch.min(in_height_coords, dim=-1) 326 | 327 | # Get left and right edges 328 | in_width, _ = torch.max(masks, dim=-2) 329 | in_width_coords = in_width * torch.arange(w, device=in_width.device)[None, :] 330 | right_edges, _ = torch.max(in_width_coords, dim=-1) 331 | in_width_coords = in_width_coords + w * (~in_width) 332 | left_edges, _ = torch.min(in_width_coords, dim=-1) 333 | 334 | # If the mask is empty the right edge will be to the left of the left edge. 335 | # Replace these boxes with [0, 0, 0, 0] 336 | empty_filter = (right_edges < left_edges) | (bottom_edges < top_edges) 337 | out = torch.stack([left_edges, top_edges, right_edges, bottom_edges], dim=-1) 338 | out = out * (~empty_filter).unsqueeze(-1) 339 | 340 | # Return to original shape 341 | if len(shape) > 2: 342 | out = out.reshape(*shape[:-2], 4) 343 | else: 344 | out = out[0] 345 | 346 | return out 347 | -------------------------------------------------------------------------------- /code/train_retraining.py: -------------------------------------------------------------------------------- 1 | import os 2 | import pdb 3 | import sys 4 | import random 5 | import shutil 6 | import argparse 7 | import logging 8 | import numpy as np 9 | import nibabel as nib 10 | from tqdm import tqdm 11 | from scipy import ndimage 12 | from importlib import import_module 13 | from tensorboardX import SummaryWriter 14 | 15 | parser = argparse.ArgumentParser() 16 | parser.add_argument('--root_path', type=str, default='../data/LA/processed_h5_rdm_4/', help='Name of Experiment') 17 | parser.add_argument('--exp', type=str, default='name', help='model_name') 18 | parser.add_argument('--dataset', type=str, default='la', help='dataset to use') 19 | parser.add_argument('--label_num', type=int, default=16, help='number of labeled data') 20 | 21 | parser.add_argument('--model_type', type=str, default='vnet', help='model_type') 22 | parser.add_argument('--rank', type=int, default=4, help='Rank for LoRA adaptation') 23 | parser.add_argument('--module', type=str, default='sam_lora_image_encoder') 24 | 25 | parser.add_argument('--patch_size', type=int, default=128, help='shape of data') 26 | parser.add_argument('--input_size', type=int, default=1024, help='shape of data') 27 | parser.add_argument('--num_classes', type=int, default=1, help='number of class') 28 | parser.add_argument('--save_img', type=int, default=250, help='img saving iterations') 29 | # load 30 | parser.add_argument('--pre_exp', type=str, default='sam_ft', help='model_name') 31 | parser.add_argument('--pre_iter', type=int, default=6000, help='maximum epoch number to train') 32 | parser.add_argument('--load', action="store_true", help='load reg & seg net') 33 | parser.add_argument('--load_iter', type=int, default=0, help='load iter') 34 | parser.add_argument('--change_lr', type=int, default=2500, help='iter for changing lr') 35 | 36 | parser.add_argument('--pre_seg_iter', type=int, default=0, help='seg number to train') 37 | parser.add_argument('--save_iter', type=int, default=1000, help='maximum epoch number to train') 38 | parser.add_argument('--max_iterations', type=int, default=6000, help='maximum epoch number to train') 39 | parser.add_argument('--batch_size', type=int, default=4, help='batch_size per gpu') 40 | parser.add_argument('--labeled_bs', type=int, default=2, help='labeled_batch_size per gpu') 41 | 42 | parser.add_argument('--lr_seg', type=float, default=0.01, help='maximum epoch number to train') 43 | parser.add_argument('--deterministic', type=int, default=1, help='whether use deterministic training') 44 | parser.add_argument('--seed', type=int, default=1337, help='random seed') 45 | parser.add_argument('--gpu', type=str, default='0', help='GPU to use') 46 | 47 | parser.add_argument('--ema_decay', type=float, default=0.99, help='ema_decay') 48 | parser.add_argument('--consistency', type=float, default=0.1, help='consistency') 49 | parser.add_argument('--consistency_rampup', type=float, default=40.0, help='consistency_rampup') 50 | args = parser.parse_args() 51 | 52 | root = "../" 53 | 54 | train_data_path = args.root_path 55 | snapshot_path = root + "model_" + args.dataset + "/" + args.exp + "/" 56 | 57 | os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu 58 | batch_size = args.batch_size 59 | labeled_bs = args.labeled_bs 60 | n_gpu = len(args.gpu.split(',')) 61 | print(batch_size) 62 | max_iterations, base_lr, input_size, patch_size = args.max_iterations, args.lr_seg, args.input_size, args.patch_size 63 | num_classes = args.num_classes 64 | 65 | import torch 66 | import torch.optim as optim 67 | import torch.nn.functional as F 68 | import torch.backends.cudnn as cudnn 69 | from torchvision import transforms 70 | from torch.utils.data import DataLoader 71 | from torch.nn.modules.loss import CrossEntropyLoss 72 | 73 | from sam_lora_image_encoder import LoRA_Sam 74 | from segment_anything_lora import sam_model_registry 75 | # from segment_anything_lora.utils.transforms import ResizeLongestSide 76 | from networks.hierarchical_vnet import VNet 77 | from networks.hierarchical_unet_3d import UNet_3D 78 | from dataloaders.dataset import * 79 | from utils import ramps, losses 80 | from utils.util import * 81 | 82 | if args.deterministic: 83 | cudnn.benchmark = False 84 | cudnn.deterministic = True 85 | random.seed(args.seed) 86 | np.random.seed(args.seed) 87 | torch.manual_seed(args.seed) 88 | torch.cuda.manual_seed(args.seed) 89 | 90 | def worker_init_fn(worker_id): 91 | random.seed(args.seed + worker_id) 92 | 93 | def update_ema_variables(model, ema_model, alpha, global_step): 94 | # Use the true average until the exponential average is more correct 95 | alpha = min(1 - 1 / (global_step + 1), alpha) 96 | for ema_param, param in zip(ema_model.parameters(), model.parameters()): 97 | ema_param.data.mul_(alpha).add_(1 - alpha, param.data) 98 | 99 | def get_current_consistency_weight(iter_num): 100 | epoch = iter_num // 150 101 | consistency_weight = args.consistency * ramps.sigmoid_rampup(epoch, args.consistency_rampup) 102 | return consistency_weight 103 | 104 | 105 | if __name__ == "__main__": 106 | ## make logger file 107 | if not os.path.exists(snapshot_path): 108 | os.makedirs(snapshot_path) 109 | os.makedirs(snapshot_path + 'saveimg') 110 | 111 | logging.basicConfig(filename=snapshot_path + "/log.txt", level=logging.INFO, 112 | format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S') 113 | logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) 114 | logging.info(str(args)) 115 | 116 | if args.dataset == 'la': 117 | db_train = LAHeart(base_dir=train_data_path, 118 | transform=transforms.Compose([ 119 | RandomRotFlip(), 120 | RandomCrop((112, 112, 80)), 121 | ToTensor(), 122 | ])) 123 | labeled_idxs = list(range(args.label_num)) 124 | unlabeled_idxs = list(range(args.label_num, 80)) 125 | 126 | elif args.dataset == 'btcv': 127 | db_train = BTCV(base_dir=train_data_path, 128 | transform=transforms.Compose([ 129 | RandomCrop((patch_size, patch_size, patch_size)), 130 | ToTensor(), 131 | ])) 132 | labeled_idxs = list(range(args.label_num)) 133 | unlabeled_idxs = list(range(args.label_num, 18)) 134 | 135 | elif args.dataset == 'mact': 136 | db_train = MACT(base_dir=train_data_path, 137 | transform=transforms.Compose([ 138 | RandomCrop((patch_size, patch_size, patch_size)), 139 | ToTensor(), 140 | ])) 141 | labeled_idxs = list(range(args.label_num)) 142 | unlabeled_idxs = list(range(args.label_num, 70)) 143 | 144 | elif args.dataset == 'brats': 145 | db_train = BraTS19(base_dir=train_data_path, 146 | transform=transforms.Compose([ 147 | RandomRotFlip(), 148 | RandomCrop((patch_size, patch_size, patch_size)), 149 | ToTensor(), 150 | ])) 151 | labeled_idxs = list(range(args.label_num)) 152 | unlabeled_idxs = list(range(args.label_num, 250)) 153 | 154 | multimask_output = True if num_classes > 2 else False 155 | batch_sampler = TwoStreamBatchSampler(labeled_idxs, unlabeled_idxs, batch_size, batch_size-labeled_bs) 156 | trainloader = DataLoader(db_train, batch_sampler=batch_sampler, num_workers=4, pin_memory=False, worker_init_fn=worker_init_fn) 157 | 158 | model_sam, img_embedding_size = sam_model_registry["vit_b"](image_size=args.input_size, num_classes=num_classes, checkpoint='pre_weight/sam_vit_b_01ec64.pth', pixel_mean=[0, 0, 0], pixel_std=[1, 1, 1]) 159 | pkg = import_module(args.module) 160 | model_sam = pkg.LoRA_Sam(model_sam, args.rank).cuda() 161 | sam_checkpoint = root + "model_" + args.dataset + "/" + args.pre_exp + "/sam_iter_" + str(args.pre_iter) + ".pth" 162 | model_sam.load_lora_parameters(sam_checkpoint) 163 | print("sam weight from {}".format(sam_checkpoint)) 164 | 165 | def create_model(ema=False): 166 | if args.model_type == "vnet": 167 | net = VNet(n_channels=1, n_classes=num_classes + 1, normalization='batchnorm', has_dropout=True, pyramid_has_dropout=True) 168 | elif args.model_type == "unet_3d": 169 | net = UNet_3D(in_channels=1, n_classes=num_classes + 1) 170 | model = net.cuda() 171 | if ema: 172 | for param in model.parameters(): 173 | param.detach_() 174 | return model 175 | 176 | model_seg = create_model() 177 | model_seg_ema = create_model(ema=True) 178 | 179 | if args.load: 180 | seg_checkpoint = os.path.join(snapshot_path, 'iter_' + str(args.load_iter) + '.pth') 181 | model_seg.load_state_dict(torch.load(seg_checkpoint)) 182 | print("init weight from {}".format(seg_checkpoint)) 183 | 184 | model_sam.eval() 185 | model_seg.train() 186 | model_seg_ema.train() 187 | 188 | # Set optimizer and losses 189 | optimizer_seg = optim.SGD(model_seg.parameters(), lr=base_lr, momentum=0.9, weight_decay=0.0001) 190 | 191 | ce_loss = CrossEntropyLoss() 192 | dice_loss = losses.DiceLoss(num_classes+1) 193 | 194 | writer = SummaryWriter(snapshot_path + '/log') 195 | logging.info("{} itertations per epoch".format(len(trainloader))) 196 | 197 | iter_num = args.load_iter 198 | max_epoch = (max_iterations - args.load_iter) // len(trainloader) + 1 199 | lr_ = base_lr 200 | kl_distance = torch.nn.KLDivLoss(reduction='none') 201 | 202 | for epoch_num in tqdm(range(max_epoch), ncols=70): 203 | for i_batch, sampled_batch in enumerate(trainloader): 204 | volume_batch, label_batch = sampled_batch['image'], sampled_batch['label'] 205 | volume_batch, label_batch = volume_batch.cuda(), label_batch.cuda() 206 | noise = torch.clamp(torch.randn_like(volume_batch) * 0.1, -0.2, 0.2) 207 | volume_batch_ema = volume_batch + noise 208 | if iter_num % args.save_img == 0: 209 | nib.save(nib.Nifti1Image(volume_batch[labeled_bs,0].cpu().detach().numpy().astype(np.float32), np.eye(4)), snapshot_path + '/saveimg' + '/img_' + str(iter_num) + '.nii.gz') 210 | nib.save(nib.Nifti1Image(label_batch[labeled_bs].cpu().detach().numpy().astype(np.float32), np.eye(4)), snapshot_path + '/saveimg' + '/gt_' + str(iter_num) + '.nii.gz') 211 | 212 | ## Train Segmentation Module 213 | outputs, _, _, _ = model_seg(volume_batch) 214 | outputs_soft = F.softmax(outputs, dim=1) 215 | prediction_seg = torch.argmax(outputs_soft, dim=1) 216 | 217 | with torch.no_grad(): 218 | outputs_ema, _, _, _ = model_seg_ema(volume_batch_ema) 219 | outputs_soft_ema = F.softmax(outputs_ema, dim=1) 220 | prediction_seg_ema = torch.argmax(outputs_soft_ema, dim=1) 221 | 222 | ## Generate pseudo labels 223 | image = volume_batch[labeled_bs:] 224 | label = label_batch[labeled_bs:] 225 | output_soft_single, prediction_sam = sam_test_tensor(model_sam, image, label, input_size, batch_size, labeled_bs, num_classes, multimask_output) 226 | 227 | # w_loss_p = 1 - iter_num / max_iterations 228 | prediction_sam = prediction_sam.long().cuda() 229 | 230 | ## calculate the loss 231 | loss_ce_seg_l = ce_loss(outputs[:labeled_bs], label_batch[:labeled_bs]) 232 | loss_dice_seg_l = dice_loss(outputs_soft[:labeled_bs], label_batch[:labeled_bs].unsqueeze(1)) 233 | loss_seg_l = 0.5 * (loss_ce_seg_l + loss_dice_seg_l) 234 | 235 | consistency_weight = get_current_consistency_weight(iter_num) 236 | 237 | loss_ce_seg_u = ce_loss(outputs[labeled_bs:], prediction_sam) 238 | loss_dice_seg_u = dice_loss(outputs_soft[labeled_bs:], prediction_sam.unsqueeze(1)) 239 | loss_seg_u = 0.5 * (loss_ce_seg_u + loss_dice_seg_u) 240 | # loss_seg_u = w_loss_p * loss_seg_u 241 | 242 | consistency_dist = losses.softmax_mse_loss(outputs, outputs_ema) 243 | consistency_dist = torch.mean(consistency_dist) 244 | consistency_loss = consistency_weight * consistency_dist 245 | 246 | loss_seg = loss_seg_l + loss_seg_u + consistency_loss 247 | 248 | optimizer_seg.zero_grad() 249 | loss_seg.backward() 250 | optimizer_seg.step() 251 | update_ema_variables(model_seg, model_seg_ema, args.ema_decay, iter_num) 252 | 253 | if iter_num % args.save_img == 0: 254 | nib.save(nib.Nifti1Image(prediction_seg[labeled_bs].cpu().detach().numpy().astype(np.float32), np.eye(4)), snapshot_path + '/saveimg' + '/pred_seg_' + str(iter_num) + '_u.nii.gz') 255 | nib.save(nib.Nifti1Image(prediction_sam[0].cpu().detach().numpy().astype(np.float32), np.eye(4)), snapshot_path + '/saveimg' + '/pred_sam_' + str(iter_num) + '_u.nii.gz') 256 | 257 | iter_num = iter_num + 1 258 | writer.add_scalar('lr', lr_, iter_num) 259 | writer.add_scalar('loss/loss_seg', loss_seg, iter_num) 260 | writer.add_scalar('loss/loss_seg_l', loss_seg_l, iter_num) 261 | writer.add_scalar('loss/loss_seg_u', loss_seg_u, iter_num) 262 | logging.info('iter %d : seg loss : %f, sup loss : %f, unsup loss : %f' % (iter_num, loss_seg, loss_seg_l, loss_seg_u.item())) 263 | 264 | ## change lr 265 | if args.dataset == "btcv": 266 | lr_ = base_lr * (1.0 - iter_num / max_iterations) ** 0.9 267 | else: 268 | if iter_num % args.change_lr == 0: 269 | lr_ = base_lr * 0.1 ** (iter_num // args.change_lr) 270 | for param_group in optimizer_seg.param_groups: 271 | param_group['lr'] = lr_ 272 | 273 | if iter_num % args.save_iter == 0: 274 | save_mode_path = os.path.join(snapshot_path, 'iter_' + str(iter_num) + '.pth') 275 | torch.save(model_seg.state_dict(), save_mode_path) 276 | logging.info("save model to {}".format(save_mode_path)) 277 | 278 | if iter_num >= max_iterations: 279 | break 280 | if iter_num >= max_iterations: 281 | break 282 | save_mode_path = os.path.join(snapshot_path, 'iter_' + str(max_iterations) + '.pth') 283 | torch.save(model_seg.state_dict(), save_mode_path) 284 | logging.info("save model to {}".format(save_mode_path)) 285 | writer.close() 286 | 287 | -------------------------------------------------------------------------------- /code/dataloaders/dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import pdb 3 | 4 | import torch 5 | import numpy as np 6 | from glob import glob 7 | from torch.utils.data import Dataset 8 | import h5py 9 | import itertools 10 | from torch.utils.data.sampler import Sampler 11 | import SimpleITK as sitk 12 | 13 | def resampling(roiImg, new_size, lbl=False): 14 | new_spacing = [old_sz * old_spc / new_sz for old_sz, old_spc, new_sz in 15 | zip(roiImg.GetSize(), roiImg.GetSpacing(), new_size)] 16 | if lbl: 17 | resampled_sitk = sitk.Resample(roiImg, new_size, sitk.Transform(), sitk.sitkNearestNeighbor, roiImg.GetOrigin(), 18 | new_spacing, roiImg.GetDirection(), 0.0, roiImg.GetPixelIDValue()) 19 | else: 20 | resampled_sitk = sitk.Resample(roiImg, new_size, sitk.Transform(), sitk.sitkLinear, roiImg.GetOrigin(), 21 | new_spacing, roiImg.GetDirection(), 0.0, roiImg.GetPixelIDValue()) 22 | return resampled_sitk 23 | 24 | 25 | class LAHeart(Dataset): 26 | """ LA Dataset """ 27 | def __init__(self, base_dir=None, list_num='', num=None, transform=None): 28 | self._base_dir = base_dir 29 | self.transform = transform 30 | self.sample_list = [] 31 | 32 | with open(self._base_dir+'/../train' + list_num + '.list', 'r') as f: 33 | self.image_list = f.readlines() 34 | self.image_list = [item.replace('\n','') for item in self.image_list] 35 | if num is not None: 36 | self.image_list = self.image_list[:num] 37 | print("total {} samples".format(len(self.image_list))) 38 | 39 | def __len__(self): 40 | return len(self.image_list) 41 | 42 | def __getitem__(self, idx): 43 | image_name = self.image_list[idx] 44 | h5f = h5py.File(self._base_dir+"/"+image_name+"/mri_norm2.h5", 'r') 45 | image, label = h5f['image'][:], h5f['label'][:] 46 | image = (image - np.mean(image)) / np.std(image) 47 | sample = {'image': image, 'label': label} 48 | if self.transform: 49 | sample = self.transform(sample) 50 | return sample 51 | 52 | 53 | class LAHeart_unlab(Dataset): 54 | """ LA Dataset """ 55 | def __init__(self, base_dir=None, list_num='', label_num=None): 56 | self._base_dir = base_dir 57 | self.sample_list = [] 58 | 59 | with open(self._base_dir+'/../train' + list_num + '.list', 'r') as f: 60 | self.image_list = f.readlines() 61 | self.image_list = [item.replace('\n','') for item in self.image_list] 62 | if label_num is not None: 63 | self.image_list = self.image_list[label_num:] 64 | print("total {} samples".format(len(self.image_list))) 65 | 66 | def __len__(self): 67 | return len(self.image_list) 68 | 69 | def __getitem__(self, idx): 70 | image_name = self.image_list[idx] 71 | h5f = h5py.File(self._base_dir+"/"+image_name+"/mri_norm2.h5", 'r') 72 | image, label = h5f['image'][:], h5f['label_full'][:] 73 | image = (image - np.mean(image)) / np.std(image) 74 | sample = {'name': image_name, 'image': image, 'label': label} 75 | # sample = {'image': image, 'label': label} 76 | return sample 77 | 78 | 79 | class BTCV(Dataset): 80 | """ BTCV Dataset """ 81 | def __init__(self, base_dir=None, num=None, transform=None): 82 | self._base_dir = base_dir 83 | self.transform = transform 84 | 85 | # with open(self._base_dir+'/../train.list', 'r') as f: 86 | with open(self._base_dir+'/../train_magic.list', 'r') as f: 87 | self.image_list = f.readlines() 88 | self.image_list = [item.replace('\n','') for item in self.image_list] 89 | if num is not None: 90 | self.image_list = self.image_list[:num] 91 | print("total {} samples".format(len(self.image_list))) 92 | 93 | def __len__(self): 94 | return len(self.image_list) 95 | 96 | def __getitem__(self, idx): 97 | image_name = self.image_list[idx] 98 | image_path = self._base_dir + '/{}.h5'.format(image_name) 99 | h5f = h5py.File(image_path, 'r') 100 | image, label = h5f['image'][:], h5f['label'][:] # (314, 314, 235) 101 | image = (image - np.mean(image)) / np.std(image) 102 | sample = {'image': image, 'label': label} 103 | if self.transform: 104 | sample = self.transform(sample) 105 | return sample 106 | 107 | 108 | class MACT(Dataset): 109 | """ MACT Dataset """ 110 | def __init__(self, base_dir=None, num=None, transform=None): 111 | self._base_dir = base_dir 112 | self.transform = transform 113 | 114 | with open(self._base_dir+'/../train.list', 'r') as f: 115 | self.image_list = f.readlines() 116 | self.image_list = [item.replace('\n','') for item in self.image_list] 117 | if num is not None: 118 | self.image_list = self.image_list[:num] 119 | print("total {} samples".format(len(self.image_list))) 120 | 121 | def __len__(self): 122 | return len(self.image_list) 123 | 124 | def __getitem__(self, idx): 125 | image_name = self.image_list[idx] 126 | image_path = self._base_dir + '/{}.h5'.format(image_name) 127 | h5f = h5py.File(image_path, 'r') 128 | image, label = h5f['image'][:], h5f['label'][:] # (314, 314, 235) 129 | image = (image - np.mean(image)) / np.std(image) 130 | sample = {'image': image, 'label': label} 131 | if self.transform: 132 | sample = self.transform(sample) 133 | return sample 134 | 135 | 136 | class BraTS19(Dataset): 137 | """ BraTS2019 Dataset """ 138 | def __init__(self, base_dir=None, num=None, transform=None): 139 | self._base_dir = base_dir 140 | self.transform = transform 141 | self.sample_list = [] 142 | 143 | with open(self._base_dir + '/../train_follow.list', 'r') as f: 144 | self.image_list = f.readlines() 145 | self.image_list = [item.replace('\n', '').split(",")[0] for item in self.image_list] 146 | if num is not None: 147 | self.image_list = self.image_list[:num] 148 | print("total {} samples".format(len(self.image_list))) 149 | 150 | def __len__(self): 151 | return len(self.image_list) 152 | 153 | def __getitem__(self, idx): 154 | image_name = self.image_list[idx] 155 | h5f = h5py.File(self._base_dir + "/{}.h5".format(image_name), 'r') 156 | image, label = h5f['image'][:], h5f['label'][:] 157 | image = image.swapaxes(0, 2) 158 | label = label.swapaxes(0, 2) 159 | image = (image - np.mean(image)) / np.std(image) 160 | label[label > 0] = 1 161 | sample = {'image': image, 'label': label.astype(np.uint8)} 162 | if self.transform: 163 | sample = self.transform(sample) 164 | return sample 165 | 166 | 167 | class BraTS19_unlab(Dataset): 168 | """ BraTS2019 Dataset """ 169 | def __init__(self, base_dir=None, label_num=None): 170 | self._base_dir = base_dir 171 | self.sample_list = [] 172 | 173 | with open(self._base_dir + '/../train_follow.list', 'r') as f: 174 | self.image_list = f.readlines() 175 | self.image_list = [item.replace('\n', '').split(",")[0] for item in self.image_list] 176 | if label_num is not None: 177 | self.image_list = self.image_list[label_num:] 178 | print("total {} samples".format(len(self.image_list))) 179 | 180 | def __len__(self): 181 | return len(self.image_list) 182 | 183 | def __getitem__(self, idx): 184 | image_name = self.image_list[idx] 185 | h5f = h5py.File(self._base_dir + "/{}.h5".format(image_name), 'r') 186 | image, label = h5f['image'][:], h5f['label'][:] 187 | image = image.swapaxes(0, 2) 188 | label = label.swapaxes(0, 2) 189 | image = (image - np.mean(image)) / np.std(image) 190 | label[label > 0] = 1 191 | sample = {'image': image, 'label': label} 192 | return sample 193 | 194 | 195 | class Resample(object): 196 | """ 197 | Crop randomly the image in a sample 198 | Args: 199 | output_size (int): Desired output size 200 | """ 201 | 202 | def __init__(self, output_size): 203 | self.output_size = output_size 204 | 205 | def __call__(self, sample): 206 | image, label = sample['image'], sample['label'] 207 | new_size = self.output_size 208 | 209 | image_itk = resampling(sitk.GetImageFromArray(image), new_size, lbl=False) 210 | label_itk = resampling(sitk.GetImageFromArray(label), new_size, lbl=True) 211 | image = sitk.GetArrayFromImage(image_itk) 212 | label = sitk.GetArrayFromImage(label_itk) 213 | 214 | return {'image': image, 'label': label} 215 | 216 | 217 | class RandomCrop(object): 218 | """ 219 | Crop randomly the image in a sample 220 | Args: 221 | output_size (int): Desired output size 222 | """ 223 | 224 | def __init__(self, output_size): 225 | self.output_size = output_size 226 | 227 | def __call__(self, sample): 228 | image, label = sample['image'], sample['label'] 229 | # pad the sample if necessary 230 | if label.shape[0] <= self.output_size[0] or label.shape[1] <= self.output_size[1] or label.shape[2] <= \ 231 | self.output_size[2]: 232 | pw = max((self.output_size[0] - label.shape[0]) // 2 + 3, 0) 233 | ph = max((self.output_size[1] - label.shape[1]) // 2 + 3, 0) 234 | pd = max((self.output_size[2] - label.shape[2]) // 2 + 3, 0) 235 | image = np.pad(image, [(pw, pw), (ph, ph), (pd, pd)], mode='constant', constant_values=0) 236 | label = np.pad(label, [(pw, pw), (ph, ph), (pd, pd)], mode='constant', constant_values=0) 237 | 238 | (w, h, d) = image.shape 239 | w1 = np.random.randint(0, w - self.output_size[0]) 240 | h1 = np.random.randint(0, h - self.output_size[1]) 241 | d1 = np.random.randint(0, d - self.output_size[2]) 242 | 243 | image = image[w1:w1 + self.output_size[0], h1:h1 + self.output_size[1], d1:d1 + self.output_size[2]] 244 | label = label[w1:w1 + self.output_size[0], h1:h1 + self.output_size[1], d1:d1 + self.output_size[2]] 245 | return {'image': image, 'label': label} 246 | 247 | 248 | class CenterCrop(object): 249 | def __init__(self, output_size): 250 | self.output_size = output_size 251 | 252 | def __call__(self, sample): 253 | image, label = sample['image'], sample['label'] 254 | 255 | # pad the sample if necessary 256 | if label.shape[0] <= self.output_size[0] or label.shape[1] <= self.output_size[1] or label.shape[2] <= \ 257 | self.output_size[2]: 258 | pw = max((self.output_size[0] - label.shape[0]) // 2 + 3, 0) 259 | ph = max((self.output_size[1] - label.shape[1]) // 2 + 3, 0) 260 | pd = max((self.output_size[2] - label.shape[2]) // 2 + 3, 0) 261 | image = np.pad(image, [(pw, pw), (ph, ph), (pd, pd)], mode='constant', constant_values=0) 262 | label = np.pad(label, [(pw, pw), (ph, ph), (pd, pd)], mode='constant', constant_values=0) 263 | 264 | (w, h, d) = image.shape 265 | 266 | w1 = int(round((w - self.output_size[0]) / 2.)) 267 | h1 = int(round((h - self.output_size[1]) / 2.)) 268 | d1 = int(round((d - self.output_size[2]) / 2.)) 269 | 270 | label = label[w1:w1 + self.output_size[0], h1:h1 + self.output_size[1], d1:d1 + self.output_size[2]] 271 | image = image[w1:w1 + self.output_size[0], h1:h1 + self.output_size[1], d1:d1 + self.output_size[2]] 272 | 273 | return {'image': image, 'label': label} 274 | 275 | class RandomRotFlip(object): 276 | """ 277 | Crop randomly flip the dataset in a sample 278 | Args: 279 | output_size (int): Desired output size 280 | """ 281 | 282 | def __call__(self, sample): 283 | image, label = sample['image'], sample['label'] 284 | k = np.random.randint(0, 4) 285 | image = np.rot90(image, k) 286 | label = np.rot90(label, k) 287 | axis = np.random.randint(0, 2) 288 | image = np.flip(image, axis=axis).copy() 289 | label = np.flip(label, axis=axis).copy() 290 | return {'image': image, 'label': label} 291 | 292 | class RandomNoise(object): 293 | def __init__(self, mu=0, sigma=0.1): 294 | self.mu = mu 295 | self.sigma = sigma 296 | 297 | def __call__(self, sample): 298 | image, label = sample['image'], sample['label'] 299 | noise = np.clip(self.sigma * np.random.randn(image.shape[0], image.shape[1], image.shape[2]), -2*self.sigma, 2*self.sigma) 300 | noise = noise + self.mu 301 | image = image + noise 302 | return {'image': image, 'label': label} 303 | 304 | 305 | class CreateOnehotLabel(object): 306 | def __init__(self, num_classes): 307 | self.num_classes = num_classes 308 | 309 | def __call__(self, sample): 310 | image, label = sample['image'], sample['label'] 311 | onehot_label = np.zeros((self.num_classes, label.shape[0], label.shape[1], label.shape[2]), dtype=np.float32) 312 | for i in range(self.num_classes): 313 | onehot_label[i, :, :, :] = (label == i).astype(np.float32) 314 | return {'image': image, 'label': label,'onehot_label':onehot_label} 315 | 316 | 317 | class ToTensor(object): 318 | """Convert ndarrays in sample to Tensors.""" 319 | 320 | def __call__(self, sample): 321 | image = sample['image'] 322 | image = image.reshape(1, image.shape[0], image.shape[1], image.shape[2]).astype(np.float32) 323 | if 'onehot_label' in sample: 324 | return {'image': torch.from_numpy(image), 'label': torch.from_numpy(sample['label']).long(), 325 | 'onehot_label': torch.from_numpy(sample['onehot_label']).long()} 326 | else: 327 | return {'image': torch.from_numpy(image), 'label': torch.from_numpy(sample['label']).long()} 328 | 329 | class TwoStreamBatchSampler(Sampler): 330 | """Iterate two sets of indices 331 | 332 | An 'epoch' is one iteration through the primary indices. 333 | During the epoch, the secondary indices are iterated through 334 | as many times as needed. 335 | """ 336 | def __init__(self, primary_indices, secondary_indices, batch_size, secondary_batch_size): 337 | self.primary_indices = primary_indices 338 | self.secondary_indices = secondary_indices 339 | self.secondary_batch_size = secondary_batch_size 340 | self.primary_batch_size = batch_size - secondary_batch_size 341 | 342 | assert len(self.primary_indices) >= self.primary_batch_size > 0 343 | assert len(self.secondary_indices) >= self.secondary_batch_size > 0 344 | 345 | def __iter__(self): 346 | primary_iter = iterate_once(self.primary_indices) 347 | secondary_iter = iterate_eternally(self.secondary_indices) 348 | return ( 349 | primary_batch + secondary_batch 350 | for (primary_batch, secondary_batch) 351 | in zip(grouper(primary_iter, self.primary_batch_size), 352 | grouper(secondary_iter, self.secondary_batch_size)) 353 | ) 354 | 355 | def __len__(self): 356 | return len(self.primary_indices) // self.primary_batch_size 357 | 358 | def iterate_once(iterable): 359 | return np.random.permutation(iterable) 360 | 361 | 362 | def iterate_eternally(indices): 363 | def infinite_shuffles(): 364 | while True: 365 | yield np.random.permutation(indices) 366 | return itertools.chain.from_iterable(infinite_shuffles()) 367 | 368 | 369 | def grouper(iterable, n): 370 | "Collect data into fixed-length chunks or blocks" 371 | # grouper('ABCDEFG', 3) --> ABC DEF" 372 | args = [iter(iterable)] * n 373 | return zip(*args) 374 | -------------------------------------------------------------------------------- /code/segment_anything_lora/modeling/image_encoder.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import torch 8 | import torch.nn as nn 9 | import torch.nn.functional as F 10 | # from icecream import ic 11 | 12 | from typing import Optional, Tuple, Type 13 | 14 | from .common import LayerNorm2d, MLPBlock 15 | 16 | 17 | # This class and its supporting functions below lightly adapted from the ViTDet backbone available at: https://github.com/facebookresearch/detectron2/blob/main/detectron2/modeling/backbone/vit.py # noqa 18 | class ImageEncoderViT(nn.Module): 19 | def __init__( 20 | self, 21 | img_size: int = 1024, 22 | patch_size: int = 16, 23 | in_chans: int = 3, 24 | embed_dim: int = 768, 25 | depth: int = 12, 26 | num_heads: int = 12, 27 | mlp_ratio: float = 4.0, 28 | out_chans: int = 256, 29 | qkv_bias: bool = True, 30 | norm_layer: Type[nn.Module] = nn.LayerNorm, 31 | act_layer: Type[nn.Module] = nn.GELU, 32 | use_abs_pos: bool = True, 33 | use_rel_pos: bool = False, 34 | rel_pos_zero_init: bool = True, 35 | window_size: int = 0, 36 | global_attn_indexes: Tuple[int, ...] = (), 37 | ) -> None: 38 | """ 39 | Args: 40 | img_size (int): Input image size. 41 | patch_size (int): Patch size. 42 | in_chans (int): Number of input image channels. 43 | embed_dim (int): Patch embedding dimension. 44 | depth (int): Depth of ViT. 45 | num_heads (int): Number of attention heads in each ViT block. 46 | mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. 47 | qkv_bias (bool): If True, add a learnable bias to query, key, value. 48 | norm_layer (nn.Module): Normalization layer. 49 | act_layer (nn.Module): Activation layer. 50 | use_abs_pos (bool): If True, use absolute positional embeddings. 51 | use_rel_pos (bool): If True, add relative positional embeddings to the attention map. 52 | rel_pos_zero_init (bool): If True, zero initialize relative positional parameters. 53 | window_size (int): Window size for window attention blocks. 54 | global_attn_indexes (list): Indexes for blocks using global attention. 55 | """ 56 | super().__init__() 57 | self.img_size = img_size 58 | 59 | self.patch_embed = PatchEmbed( 60 | kernel_size=(patch_size, patch_size), 61 | stride=(patch_size, patch_size), 62 | in_chans=in_chans, 63 | embed_dim=embed_dim, 64 | ) 65 | 66 | self.pos_embed: Optional[nn.Parameter] = None 67 | if use_abs_pos: 68 | # Initialize absolute positional embedding with pretrain image size. 69 | self.pos_embed = nn.Parameter( 70 | torch.zeros(1, img_size // patch_size, img_size // patch_size, embed_dim) 71 | ) 72 | 73 | self.blocks = nn.ModuleList() 74 | for i in range(depth): 75 | block = Block( 76 | dim=embed_dim, 77 | num_heads=num_heads, 78 | mlp_ratio=mlp_ratio, 79 | qkv_bias=qkv_bias, 80 | norm_layer=norm_layer, 81 | act_layer=act_layer, 82 | use_rel_pos=use_rel_pos, 83 | rel_pos_zero_init=rel_pos_zero_init, 84 | window_size=window_size if i not in global_attn_indexes else 0, 85 | input_size=(img_size // patch_size, img_size // patch_size), 86 | ) 87 | self.blocks.append(block) 88 | 89 | self.neck = nn.Sequential( 90 | nn.Conv2d( 91 | embed_dim, 92 | out_chans, 93 | kernel_size=1, 94 | bias=False, 95 | ), 96 | LayerNorm2d(out_chans), 97 | nn.Conv2d( 98 | out_chans, 99 | out_chans, 100 | kernel_size=3, 101 | padding=1, 102 | bias=False, 103 | ), 104 | LayerNorm2d(out_chans), 105 | ) 106 | 107 | def forward(self, x: torch.Tensor) -> torch.Tensor: 108 | x = self.patch_embed(x) # pre embed: [1, 3, 1024, 1024], post embed: [1, 64, 64, 768] 109 | if self.pos_embed is not None: 110 | x = x + self.pos_embed 111 | 112 | for blk in self.blocks: 113 | x = blk(x) 114 | 115 | x = self.neck(x.permute(0, 3, 1, 2)) # [b, c, h, w], [1, 256, 64, 64] 116 | 117 | return x 118 | 119 | 120 | class Block(nn.Module): 121 | """Transformer blocks with support of window attention and residual propagation blocks""" 122 | 123 | def __init__( 124 | self, 125 | dim: int, 126 | num_heads: int, 127 | mlp_ratio: float = 4.0, 128 | qkv_bias: bool = True, 129 | norm_layer: Type[nn.Module] = nn.LayerNorm, 130 | act_layer: Type[nn.Module] = nn.GELU, 131 | use_rel_pos: bool = False, 132 | rel_pos_zero_init: bool = True, 133 | window_size: int = 0, 134 | input_size: Optional[Tuple[int, int]] = None, 135 | ) -> None: 136 | """ 137 | Args: 138 | dim (int): Number of input channels. 139 | num_heads (int): Number of attention heads in each ViT block. 140 | mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. 141 | qkv_bias (bool): If True, add a learnable bias to query, key, value. 142 | norm_layer (nn.Module): Normalization layer. 143 | act_layer (nn.Module): Activation layer. 144 | use_rel_pos (bool): If True, add relative positional embeddings to the attention map. 145 | rel_pos_zero_init (bool): If True, zero initialize relative positional parameters. 146 | window_size (int): Window size for window attention blocks. If it equals 0, then 147 | use global attention. 148 | input_size (int or None): Input resolution for calculating the relative positional 149 | parameter size. 150 | """ 151 | super().__init__() 152 | self.norm1 = norm_layer(dim) 153 | self.attn = Attention( 154 | dim, 155 | num_heads=num_heads, 156 | qkv_bias=qkv_bias, 157 | use_rel_pos=use_rel_pos, 158 | rel_pos_zero_init=rel_pos_zero_init, 159 | input_size=input_size if window_size == 0 else (window_size, window_size), 160 | ) 161 | 162 | self.norm2 = norm_layer(dim) 163 | self.mlp = MLPBlock(embedding_dim=dim, mlp_dim=int(dim * mlp_ratio), act=act_layer) 164 | 165 | self.window_size = window_size 166 | 167 | def forward(self, x: torch.Tensor) -> torch.Tensor: 168 | shortcut = x 169 | x = self.norm1(x) 170 | # Window partition 171 | if self.window_size > 0: 172 | H, W = x.shape[1], x.shape[2] 173 | x, pad_hw = window_partition(x, self.window_size) # [B * num_windows, window_size, window_size, C] 174 | 175 | x = self.attn(x) 176 | # Reverse window partition 177 | if self.window_size > 0: 178 | x = window_unpartition(x, self.window_size, pad_hw, (H, W)) 179 | 180 | x = shortcut + x 181 | x = x + self.mlp(self.norm2(x)) 182 | 183 | return x 184 | 185 | 186 | class Attention(nn.Module): 187 | """Multi-head Attention block with relative position embeddings.""" 188 | 189 | def __init__( 190 | self, 191 | dim: int, 192 | num_heads: int = 8, 193 | qkv_bias: bool = True, 194 | use_rel_pos: bool = False, 195 | rel_pos_zero_init: bool = True, 196 | input_size: Optional[Tuple[int, int]] = None, 197 | ) -> None: 198 | """ 199 | Args: 200 | dim (int): Number of input channels. 201 | num_heads (int): Number of attention heads. 202 | qkv_bias (bool: If True, add a learnable bias to query, key, value. 203 | rel_pos (bool): If True, add relative positional embeddings to the attention map. 204 | rel_pos_zero_init (bool): If True, zero initialize relative positional parameters. 205 | input_size (int or None): Input resolution for calculating the relative positional 206 | parameter size. 207 | """ 208 | super().__init__() 209 | self.num_heads = num_heads 210 | head_dim = dim // num_heads 211 | self.scale = head_dim**-0.5 212 | 213 | self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) 214 | self.proj = nn.Linear(dim, dim) 215 | 216 | self.use_rel_pos = use_rel_pos 217 | if self.use_rel_pos: 218 | assert ( 219 | input_size is not None 220 | ), "Input size must be provided if using relative positional encoding." 221 | # initialize relative positional embeddings 222 | self.rel_pos_h = nn.Parameter(torch.zeros(2 * input_size[0] - 1, head_dim)) 223 | self.rel_pos_w = nn.Parameter(torch.zeros(2 * input_size[1] - 1, head_dim)) 224 | 225 | def forward(self, x: torch.Tensor) -> torch.Tensor: 226 | B, H, W, _ = x.shape 227 | # qkv with shape (3, B, nHead, H * W, C) 228 | qkv = self.qkv(x).reshape(B, H * W, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4) 229 | # q, k, v with shape (B * nHead, H * W, C) 230 | q, k, v = qkv.reshape(3, B * self.num_heads, H * W, -1).unbind(0) 231 | 232 | attn = (q * self.scale) @ k.transpose(-2, -1) 233 | 234 | if self.use_rel_pos: 235 | attn = add_decomposed_rel_pos(attn, q, self.rel_pos_h, self.rel_pos_w, (H, W), (H, W)) 236 | 237 | attn = attn.softmax(dim=-1) 238 | x = (attn @ v).view(B, self.num_heads, H, W, -1).permute(0, 2, 3, 1, 4).reshape(B, H, W, -1) 239 | x = self.proj(x) 240 | 241 | return x 242 | 243 | 244 | def window_partition(x: torch.Tensor, window_size: int) -> Tuple[torch.Tensor, Tuple[int, int]]: 245 | """ 246 | Partition into non-overlapping windows with padding if needed. 247 | Args: 248 | x (tensor): input tokens with [B, H, W, C]. 249 | window_size (int): window size. 250 | 251 | Returns: 252 | windows: windows after partition with [B * num_windows, window_size, window_size, C]. 253 | (Hp, Wp): padded height and width before partition 254 | """ 255 | B, H, W, C = x.shape 256 | 257 | pad_h = (window_size - H % window_size) % window_size 258 | pad_w = (window_size - W % window_size) % window_size 259 | if pad_h > 0 or pad_w > 0: 260 | x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h)) 261 | Hp, Wp = H + pad_h, W + pad_w 262 | 263 | x = x.view(B, Hp // window_size, window_size, Wp // window_size, window_size, C) 264 | windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) 265 | return windows, (Hp, Wp) 266 | 267 | 268 | def window_unpartition( 269 | windows: torch.Tensor, window_size: int, pad_hw: Tuple[int, int], hw: Tuple[int, int] 270 | ) -> torch.Tensor: 271 | """ 272 | Window unpartition into original sequences and removing padding. 273 | Args: 274 | x (tensor): input tokens with [B * num_windows, window_size, window_size, C]. 275 | window_size (int): window size. 276 | pad_hw (Tuple): padded height and width (Hp, Wp). 277 | hw (Tuple): original height and width (H, W) before padding. 278 | 279 | Returns: 280 | x: unpartitioned sequences with [B, H, W, C]. 281 | """ 282 | Hp, Wp = pad_hw 283 | H, W = hw 284 | B = windows.shape[0] // (Hp * Wp // window_size // window_size) 285 | x = windows.view(B, Hp // window_size, Wp // window_size, window_size, window_size, -1) 286 | x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, Hp, Wp, -1) 287 | 288 | if Hp > H or Wp > W: 289 | x = x[:, :H, :W, :].contiguous() 290 | return x 291 | 292 | 293 | def get_rel_pos(q_size: int, k_size: int, rel_pos: torch.Tensor) -> torch.Tensor: 294 | """ 295 | Get relative positional embeddings according to the relative positions of 296 | query and key sizes. 297 | Args: 298 | q_size (int): size of query q. 299 | k_size (int): size of key k. 300 | rel_pos (Tensor): relative position embeddings (L, C). 301 | 302 | Returns: 303 | Extracted positional embeddings according to relative positions. 304 | """ 305 | max_rel_dist = int(2 * max(q_size, k_size) - 1) 306 | # Interpolate rel pos if needed. 307 | if rel_pos.shape[0] != max_rel_dist: 308 | # Interpolate rel pos. 309 | rel_pos_resized = F.interpolate( 310 | rel_pos.reshape(1, rel_pos.shape[0], -1).permute(0, 2, 1), 311 | size=max_rel_dist, 312 | mode="linear", 313 | ) 314 | rel_pos_resized = rel_pos_resized.reshape(-1, max_rel_dist).permute(1, 0) 315 | else: 316 | rel_pos_resized = rel_pos 317 | 318 | # Scale the coords with short length if shapes for q and k are different. 319 | q_coords = torch.arange(q_size)[:, None] * max(k_size / q_size, 1.0) 320 | k_coords = torch.arange(k_size)[None, :] * max(q_size / k_size, 1.0) 321 | relative_coords = (q_coords - k_coords) + (k_size - 1) * max(q_size / k_size, 1.0) 322 | 323 | return rel_pos_resized[relative_coords.long()] 324 | 325 | 326 | def add_decomposed_rel_pos( 327 | attn: torch.Tensor, 328 | q: torch.Tensor, 329 | rel_pos_h: torch.Tensor, 330 | rel_pos_w: torch.Tensor, 331 | q_size: Tuple[int, int], 332 | k_size: Tuple[int, int], 333 | ) -> torch.Tensor: 334 | """ 335 | Calculate decomposed Relative Positional Embeddings from :paper:`mvitv2`. 336 | https://github.com/facebookresearch/mvit/blob/19786631e330df9f3622e5402b4a419a263a2c80/mvit/models/attention.py # noqa B950 337 | Args: 338 | attn (Tensor): attention map. 339 | q (Tensor): query q in the attention layer with shape (B, q_h * q_w, C). 340 | rel_pos_h (Tensor): relative position embeddings (Lh, C) for height axis. 341 | rel_pos_w (Tensor): relative position embeddings (Lw, C) for width axis. 342 | q_size (Tuple): spatial sequence size of query q with (q_h, q_w). 343 | k_size (Tuple): spatial sequence size of key k with (k_h, k_w). 344 | 345 | Returns: 346 | attn (Tensor): attention map with added relative positional embeddings. 347 | """ 348 | q_h, q_w = q_size 349 | k_h, k_w = k_size 350 | Rh = get_rel_pos(q_h, k_h, rel_pos_h) 351 | Rw = get_rel_pos(q_w, k_w, rel_pos_w) 352 | 353 | B, _, dim = q.shape 354 | r_q = q.reshape(B, q_h, q_w, dim) 355 | rel_h = torch.einsum("bhwc,hkc->bhwk", r_q, Rh) 356 | rel_w = torch.einsum("bhwc,wkc->bhwk", r_q, Rw) 357 | 358 | attn = ( 359 | attn.view(B, q_h, q_w, k_h, k_w) + rel_h[:, :, :, :, None] + rel_w[:, :, :, None, :] 360 | ).view(B, q_h * q_w, k_h * k_w) 361 | 362 | return attn 363 | 364 | 365 | class PatchEmbed(nn.Module): 366 | """ 367 | Image to Patch Embedding. 368 | """ 369 | 370 | def __init__( 371 | self, 372 | kernel_size: Tuple[int, int] = (16, 16), 373 | stride: Tuple[int, int] = (16, 16), 374 | padding: Tuple[int, int] = (0, 0), 375 | in_chans: int = 3, 376 | embed_dim: int = 768, 377 | ) -> None: 378 | """ 379 | Args: 380 | kernel_size (Tuple): kernel size of the projection layer. 381 | stride (Tuple): stride of the projection layer. 382 | padding (Tuple): padding size of the projection layer. 383 | in_chans (int): Number of input image channels. 384 | embed_dim (int): embed_dim (int): Patch embedding dimension. 385 | """ 386 | super().__init__() 387 | 388 | self.proj = nn.Conv2d( 389 | in_chans, embed_dim, kernel_size=kernel_size, stride=stride, padding=padding 390 | ) 391 | 392 | def forward(self, x: torch.Tensor) -> torch.Tensor: 393 | x = self.proj(x) 394 | # B C H W -> B H W C 395 | x = x.permute(0, 2, 3, 1) 396 | return x 397 | -------------------------------------------------------------------------------- /code/segment_anything_lora/automatic_mask_generator.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import numpy as np 8 | import torch 9 | from torchvision.ops.boxes import batched_nms, box_area # type: ignore 10 | 11 | from typing import Any, Dict, List, Optional, Tuple 12 | 13 | from .modeling import Sam 14 | from .predictor import SamPredictor 15 | from .utils.amg import ( 16 | MaskData, 17 | area_from_rle, 18 | batch_iterator, 19 | batched_mask_to_box, 20 | box_xyxy_to_xywh, 21 | build_all_layer_point_grids, 22 | calculate_stability_score, 23 | coco_encode_rle, 24 | generate_crop_boxes, 25 | is_box_near_crop_edge, 26 | mask_to_rle_pytorch, 27 | remove_small_regions, 28 | rle_to_mask, 29 | uncrop_boxes_xyxy, 30 | uncrop_masks, 31 | uncrop_points, 32 | ) 33 | 34 | 35 | class SamAutomaticMaskGenerator: 36 | def __init__( 37 | self, 38 | model: Sam, 39 | points_per_side: Optional[int] = 32, 40 | points_per_batch: int = 64, 41 | pred_iou_thresh: float = 0.88, 42 | stability_score_thresh: float = 0.95, 43 | stability_score_offset: float = 1.0, 44 | box_nms_thresh: float = 0.7, 45 | crop_n_layers: int = 0, 46 | crop_nms_thresh: float = 0.7, 47 | crop_overlap_ratio: float = 512 / 1500, 48 | crop_n_points_downscale_factor: int = 1, 49 | point_grids: Optional[List[np.ndarray]] = None, 50 | min_mask_region_area: int = 0, 51 | output_mode: str = "binary_mask", 52 | ) -> None: 53 | """ 54 | Using a SAM model, generates masks for the entire image. 55 | Generates a grid of point prompts over the image, then filters 56 | low quality and duplicate masks. The default settings are chosen 57 | for SAM with a ViT-H backbone. 58 | 59 | Arguments: 60 | model (Sam): The SAM model to use for mask prediction. 61 | points_per_side (int or None): The number of points to be sampled 62 | along one side of the image. The total number of points is 63 | points_per_side**2. If None, 'point_grids' must provide explicit 64 | point sampling. 65 | points_per_batch (int): Sets the number of points run simultaneously 66 | by the model. Higher numbers may be faster but use more GPU memory. 67 | pred_iou_thresh (float): A filtering threshold in [0,1], using the 68 | model's predicted mask quality. 69 | stability_score_thresh (float): A filtering threshold in [0,1], using 70 | the stability of the mask under changes to the cutoff used to binarize 71 | the model's mask predictions. 72 | stability_score_offset (float): The amount to shift the cutoff when 73 | calculated the stability score. 74 | box_nms_thresh (float): The box IoU cutoff used by non-maximal 75 | suppression to filter duplicate masks. 76 | crops_n_layers (int): If >0, mask prediction will be run again on 77 | crops of the image. Sets the number of layers to run, where each 78 | layer has 2**i_layer number of image crops. 79 | crops_nms_thresh (float): The box IoU cutoff used by non-maximal 80 | suppression to filter duplicate masks between different crops. 81 | crop_overlap_ratio (float): Sets the degree to which crops overlap. 82 | In the first crop layer, crops will overlap by this fraction of 83 | the image length. Later layers with more crops scale down this overlap. 84 | crop_n_points_downscale_factor (int): The number of points-per-side 85 | sampled in layer n is scaled down by crop_n_points_downscale_factor**n. 86 | point_grids (list(np.ndarray) or None): A list over explicit grids 87 | of points used for sampling, normalized to [0,1]. The nth grid in the 88 | list is used in the nth crop layer. Exclusive with points_per_side. 89 | min_mask_region_area (int): If >0, postprocessing will be applied 90 | to remove disconnected regions and holes in masks with area smaller 91 | than min_mask_region_area. Requires opencv. 92 | output_mode (str): The form masks are returned in. Can be 'binary_mask', 93 | 'uncompressed_rle', or 'coco_rle'. 'coco_rle' requires pycocotools. 94 | For large resolutions, 'binary_mask' may consume large amounts of 95 | memory. 96 | """ 97 | 98 | assert (points_per_side is None) != ( 99 | point_grids is None 100 | ), "Exactly one of points_per_side or point_grid must be provided." 101 | if points_per_side is not None: 102 | self.point_grids = build_all_layer_point_grids( 103 | points_per_side, 104 | crop_n_layers, 105 | crop_n_points_downscale_factor, 106 | ) 107 | elif point_grids is not None: 108 | self.point_grids = point_grids 109 | else: 110 | raise ValueError("Can't have both points_per_side and point_grid be None.") 111 | 112 | assert output_mode in [ 113 | "binary_mask", 114 | "uncompressed_rle", 115 | "coco_rle", 116 | ], f"Unknown output_mode {output_mode}." 117 | if output_mode == "coco_rle": 118 | from pycocotools import mask as mask_utils # type: ignore # noqa: F401 119 | 120 | if min_mask_region_area > 0: 121 | import cv2 # type: ignore # noqa: F401 122 | 123 | self.predictor = SamPredictor(model) 124 | self.points_per_batch = points_per_batch 125 | self.pred_iou_thresh = pred_iou_thresh 126 | self.stability_score_thresh = stability_score_thresh 127 | self.stability_score_offset = stability_score_offset 128 | self.box_nms_thresh = box_nms_thresh 129 | self.crop_n_layers = crop_n_layers 130 | self.crop_nms_thresh = crop_nms_thresh 131 | self.crop_overlap_ratio = crop_overlap_ratio 132 | self.crop_n_points_downscale_factor = crop_n_points_downscale_factor 133 | self.min_mask_region_area = min_mask_region_area 134 | self.output_mode = output_mode 135 | 136 | @torch.no_grad() 137 | def generate(self, image: np.ndarray) -> List[Dict[str, Any]]: 138 | """ 139 | Generates masks for the given image. 140 | 141 | Arguments: 142 | image (np.ndarray): The image to generate masks for, in HWC uint8 format. 143 | 144 | Returns: 145 | list(dict(str, any)): A list over records for masks. Each record is 146 | a dict containing the following keys: 147 | segmentation (dict(str, any) or np.ndarray): The mask. If 148 | output_mode='binary_mask', is an array of shape HW. Otherwise, 149 | is a dictionary containing the RLE. 150 | bbox (list(float)): The box around the mask, in XYWH format. 151 | area (int): The area in pixels of the mask. 152 | predicted_iou (float): The model's own prediction of the mask's 153 | quality. This is filtered by the pred_iou_thresh parameter. 154 | point_coords (list(list(float))): The point coordinates input 155 | to the model to generate this mask. 156 | stability_score (float): A measure of the mask's quality. This 157 | is filtered on using the stability_score_thresh parameter. 158 | crop_box (list(float)): The crop of the image used to generate 159 | the mask, given in XYWH format. 160 | """ 161 | 162 | # Generate masks 163 | mask_data = self._generate_masks(image) 164 | 165 | # Filter small disconnected regions and holes in masks 166 | if self.min_mask_region_area > 0: 167 | mask_data = self.postprocess_small_regions( 168 | mask_data, 169 | self.min_mask_region_area, 170 | max(self.box_nms_thresh, self.crop_nms_thresh), 171 | ) 172 | 173 | # Encode masks 174 | if self.output_mode == "coco_rle": 175 | mask_data["segmentations"] = [coco_encode_rle(rle) for rle in mask_data["rles"]] 176 | elif self.output_mode == "binary_mask": 177 | mask_data["segmentations"] = [rle_to_mask(rle) for rle in mask_data["rles"]] 178 | else: 179 | mask_data["segmentations"] = mask_data["rles"] 180 | 181 | # Write mask records 182 | curr_anns = [] 183 | for idx in range(len(mask_data["segmentations"])): 184 | ann = { 185 | "segmentation": mask_data["segmentations"][idx], 186 | "area": area_from_rle(mask_data["rles"][idx]), 187 | "bbox": box_xyxy_to_xywh(mask_data["boxes"][idx]).tolist(), 188 | "predicted_iou": mask_data["iou_preds"][idx].item(), 189 | "point_coords": [mask_data["points"][idx].tolist()], 190 | "stability_score": mask_data["stability_score"][idx].item(), 191 | "crop_box": box_xyxy_to_xywh(mask_data["crop_boxes"][idx]).tolist(), 192 | } 193 | curr_anns.append(ann) 194 | 195 | return curr_anns 196 | 197 | def _generate_masks(self, image: np.ndarray) -> MaskData: 198 | orig_size = image.shape[:2] 199 | crop_boxes, layer_idxs = generate_crop_boxes( 200 | orig_size, self.crop_n_layers, self.crop_overlap_ratio 201 | ) 202 | 203 | # Iterate over image crops 204 | data = MaskData() 205 | for crop_box, layer_idx in zip(crop_boxes, layer_idxs): 206 | crop_data = self._process_crop(image, crop_box, layer_idx, orig_size) 207 | data.cat(crop_data) 208 | 209 | # Remove duplicate masks between crops 210 | if len(crop_boxes) > 1: 211 | # Prefer masks from smaller crops 212 | scores = 1 / box_area(data["crop_boxes"]) 213 | scores = scores.to(data["boxes"].device) 214 | keep_by_nms = batched_nms( 215 | data["boxes"].float(), 216 | scores, 217 | torch.zeros(len(data["boxes"])), # categories 218 | iou_threshold=self.crop_nms_thresh, 219 | ) 220 | data.filter(keep_by_nms) 221 | 222 | data.to_numpy() 223 | return data 224 | 225 | def _process_crop( 226 | self, 227 | image: np.ndarray, 228 | crop_box: List[int], 229 | crop_layer_idx: int, 230 | orig_size: Tuple[int, ...], 231 | ) -> MaskData: 232 | # Crop the image and calculate embeddings 233 | x0, y0, x1, y1 = crop_box 234 | cropped_im = image[y0:y1, x0:x1, :] 235 | cropped_im_size = cropped_im.shape[:2] 236 | self.predictor.set_image(cropped_im) 237 | 238 | # Get points for this crop 239 | points_scale = np.array(cropped_im_size)[None, ::-1] 240 | points_for_image = self.point_grids[crop_layer_idx] * points_scale 241 | 242 | # Generate masks for this crop in batches 243 | data = MaskData() 244 | for (points,) in batch_iterator(self.points_per_batch, points_for_image): 245 | batch_data = self._process_batch(points, cropped_im_size, crop_box, orig_size) 246 | data.cat(batch_data) 247 | del batch_data 248 | self.predictor.reset_image() 249 | 250 | # Remove duplicates within this crop. 251 | keep_by_nms = batched_nms( 252 | data["boxes"].float(), 253 | data["iou_preds"], 254 | torch.zeros(len(data["boxes"])), # categories 255 | iou_threshold=self.box_nms_thresh, 256 | ) 257 | data.filter(keep_by_nms) 258 | 259 | # Return to the original image frame 260 | data["boxes"] = uncrop_boxes_xyxy(data["boxes"], crop_box) 261 | data["points"] = uncrop_points(data["points"], crop_box) 262 | data["crop_boxes"] = torch.tensor([crop_box for _ in range(len(data["rles"]))]) 263 | 264 | return data 265 | 266 | def _process_batch( 267 | self, 268 | points: np.ndarray, 269 | im_size: Tuple[int, ...], 270 | crop_box: List[int], 271 | orig_size: Tuple[int, ...], 272 | ) -> MaskData: 273 | orig_h, orig_w = orig_size 274 | 275 | # Run model on this batch 276 | transformed_points = self.predictor.transform.apply_coords(points, im_size) 277 | in_points = torch.as_tensor(transformed_points, device=self.predictor.device) 278 | in_labels = torch.ones(in_points.shape[0], dtype=torch.int, device=in_points.device) 279 | masks, iou_preds, _ = self.predictor.predict_torch( 280 | in_points[:, None, :], 281 | in_labels[:, None], 282 | multimask_output=True, 283 | return_logits=True, 284 | ) 285 | 286 | # Serialize predictions and store in MaskData 287 | data = MaskData( 288 | masks=masks.flatten(0, 1), 289 | iou_preds=iou_preds.flatten(0, 1), 290 | points=torch.as_tensor(points.repeat(masks.shape[1], axis=0)), 291 | ) 292 | del masks 293 | 294 | # Filter by predicted IoU 295 | if self.pred_iou_thresh > 0.0: 296 | keep_mask = data["iou_preds"] > self.pred_iou_thresh 297 | data.filter(keep_mask) 298 | 299 | # Calculate stability score 300 | data["stability_score"] = calculate_stability_score( 301 | data["masks"], self.predictor.model.mask_threshold, self.stability_score_offset 302 | ) 303 | if self.stability_score_thresh > 0.0: 304 | keep_mask = data["stability_score"] >= self.stability_score_thresh 305 | data.filter(keep_mask) 306 | 307 | # Threshold masks and calculate boxes 308 | data["masks"] = data["masks"] > self.predictor.model.mask_threshold 309 | data["boxes"] = batched_mask_to_box(data["masks"]) 310 | 311 | # Filter boxes that touch crop boundaries 312 | keep_mask = ~is_box_near_crop_edge(data["boxes"], crop_box, [0, 0, orig_w, orig_h]) 313 | if not torch.all(keep_mask): 314 | data.filter(keep_mask) 315 | 316 | # Compress to RLE 317 | data["masks"] = uncrop_masks(data["masks"], crop_box, orig_h, orig_w) 318 | data["rles"] = mask_to_rle_pytorch(data["masks"]) 319 | del data["masks"] 320 | 321 | return data 322 | 323 | @staticmethod 324 | def postprocess_small_regions( 325 | mask_data: MaskData, min_area: int, nms_thresh: float 326 | ) -> MaskData: 327 | """ 328 | Removes small disconnected regions and holes in masks, then reruns 329 | box NMS to remove any new duplicates. 330 | 331 | Edits mask_data in place. 332 | 333 | Requires open-cv as a dependency. 334 | """ 335 | if len(mask_data["rles"]) == 0: 336 | return mask_data 337 | 338 | # Filter small disconnected regions and holes 339 | new_masks = [] 340 | scores = [] 341 | for rle in mask_data["rles"]: 342 | mask = rle_to_mask(rle) 343 | 344 | mask, changed = remove_small_regions(mask, min_area, mode="holes") 345 | unchanged = not changed 346 | mask, changed = remove_small_regions(mask, min_area, mode="islands") 347 | unchanged = unchanged and not changed 348 | 349 | new_masks.append(torch.as_tensor(mask).unsqueeze(0)) 350 | # Give score=0 to changed masks and score=1 to unchanged masks 351 | # so NMS will prefer ones that didn't need postprocessing 352 | scores.append(float(unchanged)) 353 | 354 | # Recalculate boxes and remove any new duplicates 355 | masks = torch.cat(new_masks, dim=0) 356 | boxes = batched_mask_to_box(masks) 357 | keep_by_nms = batched_nms( 358 | boxes.float(), 359 | torch.as_tensor(scores), 360 | torch.zeros(len(boxes)), # categories 361 | iou_threshold=nms_thresh, 362 | ) 363 | 364 | # Only recalculate RLEs for masks that have changed 365 | for i_mask in keep_by_nms: 366 | if scores[i_mask] == 0.0: 367 | mask_torch = masks[i_mask].unsqueeze(0) 368 | mask_data["rles"][i_mask] = mask_to_rle_pytorch(mask_torch)[0] 369 | mask_data["boxes"][i_mask] = boxes[i_mask] # update res directly 370 | mask_data.filter(keep_by_nms) 371 | 372 | return mask_data 373 | --------------------------------------------------------------------------------