#from segment_anything import SamPredictor, sam_model_registry
from models.sam import SamPredictor, sam_model_registry
from models.sam.utils.transforms import ResizeLongestSide
from models.sam.modeling.prompt_encoder import auto_cls_emb
from models.sam.modeling.prompt_encoder import attention_fusion
from skimage.measure import label
#Scientific computing
import numpy as np
import os
#Pytorch packages
import torch
from torch import nn
import torch.optim as optim
from einops import rearrange
import torchvision
from torchvision import datasets
from tensorboardX import SummaryWriter
#Visulization
import matplotlib.pyplot as plt
from torchvision import transforms
from PIL import Image
#Others
from torch.utils.data import DataLoader, Subset
from torch.autograd import Variable
import matplotlib.pyplot as plt
import copy
from dataset_bone import MRI_dataset_multicls
import torch.nn.functional as F
from torch.nn.functional import one_hot
from pathlib import Path
from tqdm import tqdm
from losses import DiceLoss
from dsc import dice_coeff,dice_coeff_multi_class
import cv2
import monai
from utils import vis_image
import random
import cfg
args = cfg.parse_args()
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
args.if_mask_decoder_adapter=True
args.if_encoder_adapter = True
args.lr = 5e-4
args.decoder_adapt_depth = 2
args.if_warmup = True
args.initial_path = '/mnt/largeDrives/sevenTBTwo/bone_proj/codes_for_data/'
args.pretrain_weight = os.path.join('/mnt/largeDrives/sevenTBTwo/bone_proj/codes_for_data/588/fine-tune-sam/Medical-SAM-Adapter','2D-MobileSAM-onlyfusion-adapter_Bone_0107_paired_attentionpredicted','checkpoint_best.pth')
args.num_classes = 2
args.targets = 'multi_all'
def train_model(trainloader,valloader,dir_checkpoint,epochs):
# Set up model
if args.if_warmup:
b_lr = args.lr / args.warmup_period
else:
b_lr = args.lr
iter_num = 0
max_iterations = epochs * len(trainloader)
writer = SummaryWriter(dir_checkpoint + '/log')
sam = sam_model_registry["vit_t"](args,checkpoint=args.pretrain_weight,num_classes=args.num_classes)
sam.load_state_dict(torch.load(os.path.join(args.pretrain_weight)), strict = False)
print(sam)
for n, value in sam.named_parameters():
value.requires_grad = False
for n, value in sam.mask_decoder.named_parameters():
if "Adapter" in n: # only update parameters in decoder adapter
value.requires_grad = True
if 'output_hypernetworks_mlps' in n:
value.requires_grad = True
print('if image encoder adapter:',args.if_encoder_adapter)
print('if mask decoder adapter:',args.if_mask_decoder_adapter)
sam.to('cuda')
optimizer = optim.AdamW(sam.parameters(), lr=args.lr, betas=(0.9, 0.999), eps=1e-08, weight_decay=0, amsgrad=False)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5) #learning rate decay
criterion1 = monai.losses.DiceLoss(sigmoid=True, squared_pred=True, to_onehot_y=True,reduction='mean')
criterion2 = nn.CrossEntropyLoss()
pbar = tqdm(range(epochs))
val_largest_dsc = 0
last_update_epoch = 0
for epoch in pbar:
sam.train()
train_loss = 0
for i,data in enumerate(trainloader):
imgs = data['image'].cuda()
img_emb= sam.image_encoder(imgs)
alpha = random.random()
# automatic masks contaning all muscles
msks = torchvision.transforms.Resize((args.out_size,args.out_size))(data['mask'])
#print('mask unique value:',msks.unique())
msks = msks.cuda()
sparse_emb, dense_emb = sam.prompt_encoder(
points=None,
boxes=None,
masks=None,
)
pred, _ = sam.mask_decoder(
image_embeddings=img_emb,
image_pe=sam.prompt_encoder.get_dense_pe(),
sparse_prompt_embeddings=sparse_emb,
dense_prompt_embeddings=dense_emb,
multimask_output=True,
)
loss_dice = criterion1(pred,msks.float())
loss_ce = criterion2(pred,torch.squeeze(msks.long(),1))
loss = loss_dice + loss_ce
loss.backward()
optimizer.step()
optimizer.zero_grad(set_to_none=True)
if args.if_warmup and iter_num < args.warmup_period:
lr_ = args.lr * ((iter_num + 1) / args.warmup_period)
for param_group in optimizer.param_groups:
param_group['lr'] = lr_
else:
if args.if_warmup:
shift_iter = iter_num - args.warmup_period
assert shift_iter >= 0, f'Shift iter is {shift_iter}, smaller than zero'
lr_ = args.lr * (1.0 - shift_iter / max_iterations) ** 0.9 # learning rate adjustment depends on the max iterations
for param_group in optimizer.param_groups:
param_group['lr'] = lr_
train_loss += loss.item()
iter_num+=1
writer.add_scalar('info/lr', lr_, iter_num)
writer.add_scalar('info/total_loss', loss, iter_num)
writer.add_scalar('info/loss_ce', loss_ce, iter_num)
writer.add_scalar('info/loss_dice', loss_dice, iter_num)
train_loss /= (i+1)
pbar.set_description('Epoch num {}| train loss {} \n'.format(epoch,train_loss))
if epoch%2==0:
eval_loss=0
dsc = 0
sam.eval()
with torch.no_grad():
for i,data in enumerate(valloader):
imgs = data['image'].cuda()
img_emb= sam.image_encoder(imgs)
alpha = random.random()
msks = torchvision.transforms.Resize((args.out_size,args.out_size))(data['mask'])
msks = msks.cuda()
sparse_emb, dense_emb = sam.prompt_encoder(
points=None,
boxes=None,
masks=None,
)
pred, _ = sam.mask_decoder(
image_embeddings=img_emb,
image_pe=sam.prompt_encoder.get_dense_pe(),
sparse_prompt_embeddings=sparse_emb,
dense_prompt_embeddings=dense_emb,
multimask_output=True,
)
loss = criterion1(pred,msks.float()) + criterion2(pred,torch.squeeze(msks.long(),1))
eval_loss +=loss.item()
dsc_batch = dice_coeff_multi_class(pred.argmax(dim=1).cpu(), torch.squeeze(msks.long(),1).cpu().long(), 5)
dsc+=dsc_batch
eval_loss /= (i+1)
dsc /= (i+1)
writer.add_scalar('eval/loss', eval_loss, epoch)
writer.add_scalar('eval/dice', dsc, epoch)
print('Eval Epoch num {} | val loss {} | dsc {} \n'.format(epoch,eval_loss,dsc))
if dsc>val_largest_dsc:
val_largest_dsc = dsc
last_update_epoch = epoch
print('largest DSC now: {}'.format(dsc))
Path(dir_checkpoint).mkdir(parents=True,exist_ok = True)
torch.save(sam.state_dict(),dir_checkpoint + '/checkpoint_best.pth')
elif (epoch-last_update_epoch)>20:
# the network haven't been updated for 20 epochs
print('Training finished###########')
break
writer.close()
if __name__ == "__main__":
bodypart = 'hip'
dataset_name = 'Bone_0820_cls'
img_folder = args.initial_path +'2D-slices/images'
mask_folder = args.initial_path + '2D-slices/masks'
train_img_list = args.initial_path + 'datalist_body_parts/img_list_12_12_train_' + bodypart + '_annotate_paired_2dslices.txt'
val_img_list = args.initial_path + 'datalist_body_parts/img_list_12_12_val_' + bodypart + '_annotate_paired_2dslices.txt'
dir_checkpoint = '2D-MobileSAM-onlyfusion-adapter_'+dataset_name+'_attentionpredicted'
num_workers = 1
if_vis = True
epochs = 200
label_mapping = args.initial_path + 'segment_names_to_labels.pickle'
train_dataset = MRI_dataset_multicls(args,img_folder, mask_folder, train_img_list,phase='train',targets=[args.targets],delete_empty_masks='subsample',label_mapping=label_mapping,if_prompt=False)
eval_dataset = MRI_dataset_multicls(args,img_folder, mask_folder, val_img_list,phase='val',targets=[args.targets],delete_empty_masks='subsample',label_mapping=label_mapping,if_prompt=False)
trainloader = DataLoader(train_dataset, batch_size=16, shuffle=True, num_workers=num_workers)
valloader = DataLoader(eval_dataset, batch_size=16, shuffle=False, num_workers=num_workers)
train_model(trainloader,valloader,dir_checkpoint,epochs)
