import argparse
import os
import torch
import torch.nn.functional as F
import json
from segment_anything_volumetric import sam_model_registry
from network.model import SegVol
from data_process.demo_data_process import process_ct_gt
import monai.transforms as transforms
from utils.monai_inferers_utils import sliding_window_inference, generate_box, select_points, build_binary_cube, build_binary_points, logits2roi_coor
from utils.visualize import draw_result
def set_parse():
# %% set up parser
parser = argparse.ArgumentParser()
parser.add_argument("--test_mode", default=True, type=bool)
parser.add_argument("--resume", type = str, default = '')
parser.add_argument("-infer_overlap", default=0.5, type=float, help="sliding window inference overlap")
parser.add_argument("-spatial_size", default=(32, 256, 256), type=tuple)
parser.add_argument("-patch_size", default=(4, 16, 16), type=tuple)
parser.add_argument('-work_dir', type=str, default='./work_dir')
### demo
parser.add_argument('--demo_config', type=str, required=True)
parser.add_argument("--clip_ckpt", type = str, default = './config/clip')
args = parser.parse_args()
return args
def dice_score(preds, labels): # on GPU
assert preds.shape[0] == labels.shape[0], "predict & target batch size don't match\n" + str(preds.shape) + str(labels.shape)
predict = preds.view(1, -1)
target = labels.view(1, -1)
if target.shape[1] < 1e8:
predict = predict.cuda()
target = target.cuda()
predict = torch.sigmoid(predict)
predict = torch.where(predict > 0.5, 1., 0.)
tp = torch.sum(torch.mul(predict, target))
den = torch.sum(predict) + torch.sum(target) + 1
dice = 2 * tp / den
if target.shape[1] < 1e8:
predict = predict.cpu()
target = target.cpu()
return dice
def zoom_in_zoom_out(args, segvol_model, image, image_resize, gt3D, gt3D_resize, categories=None):
logits_labels_record = {}
image_single_resize = image_resize
image_single = image[0,0]
ori_shape = image_single.shape
for item_idx in range(len(categories)):
# get label to generate prompts
label_single = gt3D[0][item_idx]
label_single_resize = gt3D_resize[0][item_idx]
# skip meaningless categories
if torch.sum(label_single) == 0:
print('No object, skip')
continue
# generate prompts
text_single = categories[item_idx] if args.use_text_prompt else None
if categories is not None: print(f'inference |{categories[item_idx]}| target...')
points_single = None
box_single = None
if args.use_point_prompt:
point, point_label = select_points(label_single_resize, num_positive_extra=3, num_negative_extra=3)
points_single = (point.unsqueeze(0).float().cuda(), point_label.unsqueeze(0).float().cuda())
binary_points_resize = build_binary_points(point, point_label, label_single_resize.shape)
if args.use_box_prompt:
box_single = generate_box(label_single_resize).unsqueeze(0).float().cuda()
binary_cube_resize = build_binary_cube(box_single, binary_cube_shape=label_single_resize.shape)
####################
# zoom-out inference:
print('--- zoom out inference ---')
print(f'use text-prompt [{text_single!=None}], use box-prompt [{box_single!=None}], use point-prompt [{points_single!=None}]')
with torch.no_grad():
logits_global_single = segvol_model(image_single_resize.cuda(),
text=text_single,
boxes=box_single,
points=points_single)
# resize back global logits
logits_global_single = F.interpolate(
logits_global_single.cpu(),
size=ori_shape, mode='nearest')[0][0]
# build prompt reflection for zoom-in
if args.use_point_prompt:
binary_points = F.interpolate(
binary_points_resize.unsqueeze(0).unsqueeze(0).float(),
size=ori_shape, mode='nearest')[0][0]
if args.use_box_prompt:
binary_cube = F.interpolate(
binary_cube_resize.unsqueeze(0).unsqueeze(0).float(),
size=ori_shape, mode='nearest')[0][0]
zoom_out_dice = dice_score(logits_global_single.squeeze(), label_single.squeeze())
logits_labels_record[categories[item_idx]] = (
zoom_out_dice,
image_single,
points_single,
box_single,
logits_global_single,
label_single)
print(f'zoom out inference done with zoom_out_dice: {zoom_out_dice:.4f}')
if not args.use_zoom_in:
continue
####################
# zoom-in inference:
min_d, min_h, min_w, max_d, max_h, max_w = logits2roi_coor(args.spatial_size, logits_global_single)
if min_d is None:
print('Fail to detect foreground!')
continue
# Crop roi
image_single_cropped = image_single[min_d:max_d+1, min_h:max_h+1, min_w:max_w+1].unsqueeze(0).unsqueeze(0)
global_preds = (torch.sigmoid(logits_global_single[min_d:max_d+1, min_h:max_h+1, min_w:max_w+1])>0.5).long()
assert not (args.use_box_prompt and args.use_point_prompt)
# label_single_cropped = label_single[min_d:max_d+1, min_h:max_h+1, min_w:max_w+1].unsqueeze(0).unsqueeze(0)
prompt_reflection = None
if args.use_box_prompt:
binary_cube_cropped = binary_cube[min_d:max_d+1, min_h:max_h+1, min_w:max_w+1]
prompt_reflection = (
binary_cube_cropped.unsqueeze(0).unsqueeze(0),
global_preds.unsqueeze(0).unsqueeze(0)
)
if args.use_point_prompt:
binary_points_cropped = binary_points[min_d:max_d+1, min_h:max_h+1, min_w:max_w+1]
prompt_reflection = (
binary_points_cropped.unsqueeze(0).unsqueeze(0),
global_preds.unsqueeze(0).unsqueeze(0)
)
## inference
with torch.no_grad():
logits_single_cropped = sliding_window_inference(
image_single_cropped.cuda(), prompt_reflection,
args.spatial_size, 1, segvol_model, args.infer_overlap,
text=text_single,
use_box=args.use_box_prompt,
use_point=args.use_point_prompt,
)
logits_single_cropped = logits_single_cropped.cpu().squeeze()
logits_global_single[min_d:max_d+1, min_h:max_h+1, min_w:max_w+1] = logits_single_cropped
zoom_in_dice = dice_score(logits_global_single.squeeze(), label_single.squeeze())
logits_labels_record[categories[item_idx]] = (
zoom_in_dice,
image_single,
points_single,
box_single,
logits_global_single,
label_single)
print(f'===> zoom out dice {zoom_out_dice:.4f} -> zoom-out-zoom-in dice {zoom_in_dice:.4f} <===')
return logits_labels_record
def inference_single_ct(args, segvol_model, data_item, categories):
segvol_model.eval()
image, gt3D = data_item["image"].float(), data_item["label"]
image_zoom_out, gt3D__zoom_out = data_item["zoom_out_image"].float(), data_item['zoom_out_label']
logits_labels_record = zoom_in_zoom_out(
args, segvol_model,
image.unsqueeze(0), image_zoom_out.unsqueeze(0),
gt3D.unsqueeze(0), gt3D__zoom_out.unsqueeze(0), # add batch dim
categories=categories)
# visualize
if args.visualize:
for target, values in logits_labels_record.items():
dice_score, image, point_prompt, box_prompt, logits, labels = values
print(f'{target} result with Dice score {dice_score:.4f} visualizing')
draw_result(target + f"-Dice {dice_score:.4f}", image, box_prompt, point_prompt, logits, labels, args.spatial_size, args.work_dir)
def main(args):
gpu = 0
torch.cuda.set_device(gpu)
# build model
sam_model = sam_model_registry['vit'](args=args)
segvol_model = SegVol(
image_encoder=sam_model.image_encoder,
mask_decoder=sam_model.mask_decoder,
prompt_encoder=sam_model.prompt_encoder,
clip_ckpt=args.clip_ckpt,
roi_size=args.spatial_size,
patch_size=args.patch_size,
test_mode=args.test_mode,
).cuda()
segvol_model = torch.nn.DataParallel(segvol_model, device_ids=[gpu])
# load param
if os.path.isfile(args.resume):
## Map model to be loaded to specified single GPU
loc = 'cuda:{}'.format(gpu)
checkpoint = torch.load(args.resume, map_location=loc)
segvol_model.load_state_dict(checkpoint['model'], strict=True)
print("loaded checkpoint '{}' (epoch {})".format(args.resume, checkpoint['epoch']))
# load demo config
with open(args.demo_config, 'r') as file:
config_dict = json.load(file)
ct_path, gt_path, categories = config_dict['demo_case']['ct_path'], config_dict['demo_case']['gt_path'], config_dict['categories']
# preprocess for data
data_item = process_ct_gt(ct_path, gt_path, categories, args.spatial_size) # keys: image, label
# seg config for prompt & zoom-in-zoom-out
args.use_zoom_in = True
args.use_text_prompt = True
args.use_box_prompt = True
args.use_point_prompt = False
args.visualize = False
inference_single_ct(args, segvol_model, data_item, categories)
if __name__ == "__main__":
args = set_parse()
main(args)